From 44685c496170598fe2512631e0fee042ce6e30bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Timo=20K=C3=B6ssler?= Date: Thu, 2 Jul 2026 16:54:22 +0200 Subject: [PATCH] Add data regions support --- aikido_zen/helpers/token.py | 19 +++++++ aikido_zen/helpers/token_test.py | 31 ++++++++++- aikido_zen/helpers/urls/get_api_url.py | 14 ++++- aikido_zen/helpers/urls/get_api_url_test.py | 60 +++++++++++++++++++++ 4 files changed, 121 insertions(+), 3 deletions(-) diff --git a/aikido_zen/helpers/token.py b/aikido_zen/helpers/token.py index ab39c1c14..3bd615422 100644 --- a/aikido_zen/helpers/token.py +++ b/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" diff --git a/aikido_zen/helpers/token_test.py b/aikido_zen/helpers/token_test.py index 1337a07cf..9eee9cda0 100644 --- a/aikido_zen/helpers/token_test.py +++ b/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" diff --git a/aikido_zen/helpers/urls/get_api_url.py b/aikido_zen/helpers/urls/get_api_url.py index 93bb38475..96cf99801 100644 --- a/aikido_zen/helpers/urls/get_api_url.py +++ b/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) diff --git a/aikido_zen/helpers/urls/get_api_url_test.py b/aikido_zen/helpers/urls/get_api_url_test.py index 9a8de955d..59beb6f93 100644 --- a/aikido_zen/helpers/urls/get_api_url_test.py +++ b/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