from_env_var_v2 optional args

Author: anitarua

src/momento/auth/credential_provider.py

@@ -46,7 +46,11 @@ def from_environment_variable(
         Returns:
             CredentialProvider
         """
-        warn("from_environment_variable is deprecated, use from_env_var_v2 instead", DeprecationWarning, stacklevel=2)
+        warn(
+            "from_environment_variable is deprecated, use from_environment_variables_v2 instead",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         api_key = os.getenv(env_var_name)
         if not api_key:
             raise RuntimeError(f"Missing required environment variable {env_var_name}")
@@ -115,10 +119,10 @@ def get_auth_token(self) -> str:
 
     @staticmethod
     def from_api_key_v2(api_key: str, endpoint: str) -> CredentialProvider:
-        """Creates a CredentialProvider from a global API key and endpoint.
+        """Creates a CredentialProvider from a v2 API key and endpoint.
 
         Args:
-            api_key (str): The global API key.
+            api_key (str): The v2 API key.
             endpoint (str): The Momento service endpoint.
 
         Returns:
@@ -143,12 +147,14 @@ def from_api_key_v2(api_key: str, endpoint: str) -> CredentialProvider:
         )
 
     @staticmethod
-    def from_env_var_v2(api_key_env_var: str, endpoint_env_var: str) -> CredentialProvider:
-        """Creates a CredentialProvider from an endpoint and a global API key stored in an environment variable.
+    def from_environment_variables_v2(
+        api_key_env_var: str = "MOMENTO_API_KEY", endpoint_env_var: str = "MOMENTO_ENDPOINT"
+    ) -> CredentialProvider:
+        """Creates a CredentialProvider from an endpoint and v2 API key stored in the environment variables MOMENTO_API_KEY and MOMENTO_ENDPOINT.
 
         Args:
-            api_key_env_var (str): Name of the environment variable from which the global API key will be read.
-            endpoint_env_var (str): Name of the environment variable from which the Momento service endpoint will be read.
+            api_key_env_var (str): Optionally provide an alternate environment variable name from which the v2 API key will be read.
+            endpoint_env_var (str): Optionally provide an alternate environment variable name from which the Momento service endpoint will be read.
 
         Returns:
             CredentialProvider

src/momento/auth/momento_endpoint_resolver.py

@@ -48,7 +48,7 @@ def resolve(auth_token: str) -> _TokenAndEndpoints:
     else:
         if _is_v2_api_key(auth_token):
             raise InvalidArgumentException(
-                "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_env_var_v2()` instead?",
+                "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_environment_variables_v2()` instead?",
                 Service.AUTH,
             )
         return _get_endpoint_from_token(auth_token)

tests/momento/auth/test_credential_provider.py

@@ -26,11 +26,11 @@
 os.environ[test_v1_env_var_name] = test_encoded_v1_token.decode("utf-8")
 
 # For v2 API key tests
-test_v2_key_message = {"t": "g", "id": "some-id"}
+test_v2_key_message = {"t": "g", "jti": "some-id"}
 test_v2_api_key = jwt.encode(test_v2_key_message, "secret", algorithm="HS512")
-test_v2_key_env_var_name = uuid_str()
+test_v2_key_env_var_name = "MOMENTO_API_KEY"
 test_v2_endpoint = "testEndpoint"
-test_v2_endpoint_env_var_name = uuid_str()
+test_v2_endpoint_env_var_name = "MOMENTO_ENDPOINT"
 os.environ[test_v2_key_env_var_name] = test_v2_api_key
 os.environ[test_v2_endpoint_env_var_name] = test_v2_endpoint
 
@@ -113,7 +113,6 @@ def test_env_token_raises_if_not_exists() -> None:
 @pytest.mark.parametrize(
     "provider, expected_api_key, expected_control_endpoint, expected_cache_endpoint, expected_token_endpoint",
     [
-        # v2_key_from_string - basic usage
         (
             CredentialProvider.from_api_key_v2(
                 api_key=test_v2_api_key,
@@ -124,9 +123,8 @@ def test_env_token_raises_if_not_exists() -> None:
             f"cache.{test_v2_endpoint}",
             f"token.{test_v2_endpoint}",
         ),
-        # v2_key_from_environment_variable - basic usage
         (
-            CredentialProvider.from_env_var_v2(
+            CredentialProvider.from_environment_variables_v2(
                 api_key_env_var=test_v2_key_env_var_name,
                 endpoint_env_var=test_v2_endpoint_env_var_name,
             ),
@@ -135,6 +133,13 @@ def test_env_token_raises_if_not_exists() -> None:
             f"cache.{test_v2_endpoint}",
             f"token.{test_v2_endpoint}",
         ),
+        (
+            CredentialProvider.from_environment_variables_v2(),
+            test_v2_api_key,
+            f"control.{test_v2_endpoint}",
+            f"cache.{test_v2_endpoint}",
+            f"token.{test_v2_endpoint}",
+        ),
     ],
 )
 def test_v2_api_key_endpoints(
@@ -162,24 +167,28 @@ def test_v2_key_from_string_raises_if_endpoint_empty() -> None:
 
 def test_v2_key_from_env_raises_if_env_var_name_empty() -> None:
     with pytest.raises(InvalidArgumentException, match="API key environment variable name cannot be empty"):
-        CredentialProvider.from_env_var_v2(api_key_env_var="", endpoint_env_var=test_v2_endpoint_env_var_name)
+        CredentialProvider.from_environment_variables_v2(
+            api_key_env_var="", endpoint_env_var=test_v2_endpoint_env_var_name
+        )
 
 
 def test_v2_key_from_env_raises_if_env_var_missing() -> None:
     with pytest.raises(RuntimeError, match="Missing required environment variable"):
-        CredentialProvider.from_env_var_v2(api_key_env_var=uuid_str(), endpoint_env_var=test_v2_endpoint_env_var_name)
+        CredentialProvider.from_environment_variables_v2(
+            api_key_env_var=uuid_str(), endpoint_env_var=test_v2_endpoint_env_var_name
+        )
 
 
 def test_v2_key_from_env_raises_if_endpoint_empty() -> None:
     with pytest.raises(InvalidArgumentException, match="Endpoint environment variable name cannot be empty"):
-        CredentialProvider.from_env_var_v2(api_key_env_var=test_v2_key_env_var_name, endpoint_env_var="")
+        CredentialProvider.from_environment_variables_v2(api_key_env_var=test_v2_key_env_var_name, endpoint_env_var="")
 
 
 def test_v2_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="Missing required environment variable"):
-        CredentialProvider.from_env_var_v2(
+        CredentialProvider.from_environment_variables_v2(
             api_key_env_var=empty_api_key_env_var, endpoint_env_var=test_v2_endpoint_env_var_name
         )
 
@@ -203,7 +212,7 @@ def test_v2_key_from_env_raises_if_base64_api_key() -> None:
             "Received an invalid v2 API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` with a legacy key instead?"
         ),
     ):
-        CredentialProvider.from_env_var_v2(
+        CredentialProvider.from_environment_variables_v2(
             api_key_env_var=test_v1_env_var_name, endpoint_env_var=test_v2_endpoint_env_var_name
         )
 
@@ -225,7 +234,7 @@ def test_v2_key_from_env_raises_if_pre_v1_token() -> None:
             "Received an invalid v2 API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` with a legacy key instead?"
         ),
     ):
-        CredentialProvider.from_env_var_v2(
+        CredentialProvider.from_environment_variables_v2(
             api_key_env_var=test_env_var_name, endpoint_env_var=test_v2_endpoint_env_var_name
         )
 
@@ -234,7 +243,7 @@ def test_v2_key_provided_to_from_string() -> None:
     with pytest.raises(
         InvalidArgumentException,
         match=re.escape(
-            "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_env_var_v2()` instead?"
+            "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_environment_variables_v2()` instead?"
         ),
     ):
         CredentialProvider.from_string(auth_token=test_v2_api_key)
@@ -244,7 +253,7 @@ def test_v2_key_provided_to_from_disposable_token() -> None:
     with pytest.raises(
         InvalidArgumentException,
         match=re.escape(
-            "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_env_var_v2()` instead?"
+            "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_environment_variables_v2()` instead?"
         ),
     ):
         CredentialProvider.from_disposable_token(auth_token=test_v2_api_key)