From 35d0c85bf9f0b432b45004078d1a58dbb8b60763 Mon Sep 17 00:00:00 2001 From: anitarua Date: Wed, 19 Nov 2025 16:50:05 -0800 Subject: [PATCH] feat: new credential provider methods for accepting global api keys --- src/momento/auth/credential_provider.py | 41 ++++++++++ .../momento/auth/test_credential_provider.py | 81 +++++++++++++++++++ 2 files changed, 122 insertions(+) diff --git a/src/momento/auth/credential_provider.py b/src/momento/auth/credential_provider.py index eb16000f..49d3d7be 100644 --- a/src/momento/auth/credential_provider.py +++ b/src/momento/auth/credential_provider.py @@ -102,3 +102,44 @@ def _obscure(self, value: str) -> str: def get_auth_token(self) -> str: return self.auth_token + + @staticmethod + def global_key_from_string(api_key: str, endpoint: str) -> CredentialProvider: + """Creates a CredentialProvider from a global API key and endpoint. + + Args: + api_key (str): The global API key. + endpoint (str): The Momento service endpoint. + + Returns: + CredentialProvider + """ + if len(api_key) == 0: + raise RuntimeError("API key cannot be empty.") + if len(endpoint) == 0: + raise RuntimeError("Endpoint cannot be empty.") + return CredentialProvider( + auth_token=api_key, + control_endpoint=momento_endpoint_resolver._MOMENTO_CONTROL_ENDPOINT_PREFIX + endpoint, + cache_endpoint=momento_endpoint_resolver._MOMENTO_CACHE_ENDPOINT_PREFIX + endpoint, + token_endpoint=momento_endpoint_resolver._MOMENTO_TOKEN_ENDPOINT_PREFIX + endpoint, + port=443, + ) + + @staticmethod + def global_key_from_environment_variable(env_var_name: str, endpoint: str) -> CredentialProvider: + """Creates a CredentialProvider from an endpoint and a global API key stored in an environment variable. + + Args: + env_var_name (str): Name of the environment variable from which the global API key will be read. + endpoint (str): The Momento service endpoint. + + Returns: + CredentialProvider + """ + if len(env_var_name) == 0: + raise RuntimeError("Environment variable name cannot be empty.") + api_key = os.getenv(env_var_name) + if not api_key: + raise RuntimeError(f"Missing required environment variable {env_var_name}") + return CredentialProvider.global_key_from_string(api_key, endpoint) diff --git a/tests/momento/auth/test_credential_provider.py b/tests/momento/auth/test_credential_provider.py index 2258e584..855c52b9 100644 --- a/tests/momento/auth/test_credential_provider.py +++ b/tests/momento/auth/test_credential_provider.py @@ -97,3 +97,84 @@ def test_endpoints(provider: CredentialProvider, auth_token: str, control_endpoi def test_env_token_raises_if_not_exists() -> None: with pytest.raises(RuntimeError, match=r"Missing required environment variable"): CredentialProvider.from_environment_variable(env_var_name=uuid_str()) + + +# Global API Key Tests +test_global_api_key = "testToken" +test_global_endpoint = "testEndpoint" +test_global_env_var_name = "MOMENTO_TEST_GLOBAL_API_KEY" +os.environ[test_global_env_var_name] = test_global_api_key + + +@pytest.mark.parametrize( + "provider, expected_api_key, expected_control_endpoint, expected_cache_endpoint, expected_token_endpoint", + [ + # global_key_from_string - basic usage + ( + CredentialProvider.global_key_from_string( + api_key=test_global_api_key, + endpoint=test_global_endpoint, + ), + test_global_api_key, + f"control.{test_global_endpoint}", + f"cache.{test_global_endpoint}", + f"token.{test_global_endpoint}", + ), + # global_key_from_environment_variable - basic usage + ( + CredentialProvider.global_key_from_environment_variable( + env_var_name=test_global_env_var_name, + endpoint=test_global_endpoint, + ), + test_global_api_key, + f"control.{test_global_endpoint}", + f"cache.{test_global_endpoint}", + f"token.{test_global_endpoint}", + ), + ], +) +def test_global_api_key_endpoints( + provider: CredentialProvider, + expected_api_key: str, + expected_control_endpoint: str, + expected_cache_endpoint: str, + expected_token_endpoint: str, +) -> None: + assert provider.auth_token == expected_api_key + assert provider.control_endpoint == expected_control_endpoint + assert provider.cache_endpoint == expected_cache_endpoint + assert provider.token_endpoint == expected_token_endpoint + + +def test_global_key_from_string_raises_if_api_key_empty() -> None: + with pytest.raises(RuntimeError, match=r"API key cannot be empty"): + CredentialProvider.global_key_from_string(api_key="", endpoint=test_global_endpoint) + + +def test_global_key_from_string_raises_if_endpoint_empty() -> None: + with pytest.raises(RuntimeError, match=r"Endpoint cannot be empty"): + CredentialProvider.global_key_from_string(api_key=test_global_api_key, endpoint="") + + +def test_global_key_from_env_raises_if_env_var_name_empty() -> None: + with pytest.raises(RuntimeError, match=r"Environment variable name cannot be empty"): + CredentialProvider.global_key_from_environment_variable(env_var_name="", endpoint=test_global_endpoint) + + +def test_global_key_from_env_raises_if_env_var_missing() -> None: + with pytest.raises(RuntimeError, match=r"Missing required environment variable"): + CredentialProvider.global_key_from_environment_variable(env_var_name=uuid_str(), endpoint=test_global_endpoint) + + +def test_global_key_from_env_raises_if_endpoint_empty() -> None: + with pytest.raises(RuntimeError, match=r"Endpoint cannot be empty"): + CredentialProvider.global_key_from_environment_variable(env_var_name=test_global_env_var_name, endpoint="") + + +def test_global_key_from_env_raises_if_api_key_empty_string() -> None: + empty_api_key_env_var = uuid_str() + os.environ[empty_api_key_env_var] = "" + with pytest.raises(RuntimeError, match=r"Missing required environment variable"): + CredentialProvider.global_key_from_environment_variable( + env_var_name=empty_api_key_env_var, endpoint=test_global_endpoint + )