Add data regions support

Author: timokoessler

aikido_zen/helpers/token.py

@@ -27,3 +27,22 @@ def get_token_from_env():
     if aikido_token_env is not None:
         return Token(aikido_token_env)
     return None
+
+
+def extract_region_from_token(token):
+    """
+    Extracts the region from a runtime token.
+
+    New format: AIK_RUNTIME_{sys_group_id}_{service_id}_{region}_{random}
+    Old format: AIK_RUNTIME_{sys_group_id}_{service_id}_{random}
+    """
+    if not token or not token.startswith("AIK_RUNTIME_"):
+        return "EU"
+
+    token_without_prefix = token[len("AIK_RUNTIME_") :]
+    parts = token_without_prefix.split("_")
+
+    if len(parts) == 4:
+        return parts[2]
+
+    return "EU"

aikido_zen/helpers/token_test.py

@@ -1,5 +1,5 @@
 import pytest
-from aikido_zen.helpers.token import Token
+from aikido_zen.helpers.token import Token, extract_region_from_token
 
 
 # Test Token Class :
@@ -23,3 +23,32 @@ def test_token_instance():
     token_str = "my_token"
     token = Token(token_str)
     assert isinstance(token, Token)
+
+
+# Test extract_region_from_token :
+def test_extract_region_from_token_empty():
+    assert extract_region_from_token("") == "EU"
+
+
+def test_extract_region_from_token_none():
+    assert extract_region_from_token(None) == "EU"
+
+
+def test_extract_region_from_token_no_prefix():
+    assert extract_region_from_token("my_token") == "EU"
+
+
+def test_extract_region_from_token_old_format():
+    assert extract_region_from_token("AIK_RUNTIME_123_456_random") == "EU"
+
+
+def test_extract_region_from_token_new_format():
+    assert extract_region_from_token("AIK_RUNTIME_123_456_US_random") == "US"
+
+
+def test_extract_region_from_token_new_format_me():
+    assert extract_region_from_token("AIK_RUNTIME_123_456_ME_random") == "ME"
+
+
+def test_extract_region_from_token_new_format_au():
+    assert extract_region_from_token("AIK_RUNTIME_123_456_AU_random") == "AU"

aikido_zen/helpers/urls/get_api_url.py

@@ -3,12 +3,19 @@
 import os
 from urllib.parse import urlunparse
 from aikido_zen.helpers.try_parse_url import try_parse_url
+from aikido_zen.helpers.token import extract_region_from_token, get_token_from_env
 
 DEFAULT_API_URL = "https://guard.aikido.dev/"
+REGION_API_URLS = {
+    "US": "https://guard.us.aikido.dev/",
+    "ME": "https://guard.me.aikido.dev/",
+    "AU": "https://guard.au.aikido.dev/",
+}
 
 
 def get_api_url():
-    """Checks environment var AIKIDO_ENDPOINT for the api URL"""
+    """Checks environment var AIKIDO_ENDPOINT for the api URL, falling back
+    to a region-specific URL based on the AIKIDO_TOKEN"""
     realtime_url = os.getenv("AIKIDO_ENDPOINT")
     if realtime_url is not None:
         parsed_url = try_parse_url(realtime_url)
@@ -18,4 +25,7 @@ def get_api_url():
             # Make sure ends with a slash :
             return urlunparse(parsed_url) + "/"
         return urlunparse(parsed_url)
-    return DEFAULT_API_URL
+
+    token = get_token_from_env()
+    region = extract_region_from_token(str(token) if token else "")
+    return REGION_API_URLS.get(region, DEFAULT_API_URL)

aikido_zen/helpers/urls/get_api_url_test.py

@@ -33,4 +33,64 @@ def test_get_api_url_valid_url_no_trailing_slash():
         assert get_api_url() == valid_url_no_slash + "/"
 
 
+def test_get_api_url_no_token_defaults_to_eu():
+    """Test when neither AIKIDO_ENDPOINT nor AIKIDO_TOKEN is set."""
+    with patch.dict(os.environ, {}, clear=True):
+        assert get_api_url() == DEFAULT_API_URL
+
+
+def test_get_api_url_old_format_token_defaults_to_eu():
+    """Test when AIKIDO_TOKEN is set with the old token format."""
+    with patch.dict(
+        os.environ,
+        {"AIKIDO_TOKEN": "AIK_RUNTIME_123_456_random"},
+        clear=True,
+    ):
+        assert get_api_url() == DEFAULT_API_URL
+
+
+def test_get_api_url_us_region_token():
+    """Test when AIKIDO_TOKEN encodes the US region."""
+    with patch.dict(
+        os.environ,
+        {"AIKIDO_TOKEN": "AIK_RUNTIME_123_456_US_random"},
+        clear=True,
+    ):
+        assert get_api_url() == "https://guard.us.aikido.dev/"
+
+
+def test_get_api_url_me_region_token():
+    """Test when AIKIDO_TOKEN encodes the ME region."""
+    with patch.dict(
+        os.environ,
+        {"AIKIDO_TOKEN": "AIK_RUNTIME_123_456_ME_random"},
+        clear=True,
+    ):
+        assert get_api_url() == "https://guard.me.aikido.dev/"
+
+
+def test_get_api_url_au_region_token():
+    """Test when AIKIDO_TOKEN encodes the AU region."""
+    with patch.dict(
+        os.environ,
+        {"AIKIDO_TOKEN": "AIK_RUNTIME_123_456_AU_random"},
+        clear=True,
+    ):
+        assert get_api_url() == "https://guard.au.aikido.dev/"
+
+
+def test_get_api_url_endpoint_env_var_takes_precedence_over_token():
+    """Test that AIKIDO_ENDPOINT overrides region derived from AIKIDO_TOKEN."""
+    valid_url = "https://example.com/api"
+    with patch.dict(
+        os.environ,
+        {
+            "AIKIDO_ENDPOINT": valid_url,
+            "AIKIDO_TOKEN": "AIK_RUNTIME_123_456_US_random",
+        },
+        clear=True,
+    ):
+        assert get_api_url() == valid_url + "/"
+
+
 # You may need to mock the try_parse_url function if it has side effects or is complex