Refactor Sources into a class/Protocol

Author: alextwoods

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsAuthIntegration.java

@@ -59,7 +59,7 @@ public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
                                 .nullable(true)
                                 .initialize(writer -> {
                                     writer.addImport("smithy_aws_core.credentials_resolvers", "CredentialsResolverChain");
-                                    writer.write("self.aws_credentials_identity_resolver = aws_credentials_identity_resolver or CredentialsResolverChain()");
+                                    writer.write("self.aws_credentials_identity_resolver = aws_credentials_identity_resolver or CredentialsResolverChain(config=self)");
                                 })
                                 .build())
                         .addConfigProperty(REGION)

packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/chain.py

@@ -1,27 +1,24 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
-import os
-from collections.abc import Callable, Sequence
+from collections.abc import Sequence
 
 from smithy_core.aio.interfaces.identity import IdentityResolver
 from smithy_core.exceptions import SmithyIdentityException
 from smithy_core.interfaces.identity import IdentityProperties
 
-from smithy_aws_core.credentials_resolvers import EnvironmentCredentialsResolver
+from smithy_aws_core.credentials_resolvers.environment import (
+    EnvironmentCredentialsSource,
+)
+from smithy_aws_core.credentials_resolvers.imds import IMDSCredentialsSource
+from smithy_aws_core.credentials_resolvers.interfaces import (
+    AwsCredentialsConfig,
+    CredentialsSource,
+)
 from smithy_aws_core.identity import AWSCredentialsIdentity, AWSCredentialsResolver
 
-
-def _env_creds_available() -> bool:
-    return "AWS_ACCESS_KEY_ID" in os.environ and "AWS_SECRET_ACCESS_KEY" in os.environ
-
-
-def _build_env_creds() -> AWSCredentialsResolver:
-    return EnvironmentCredentialsResolver()
-
-
-type CredentialSource = tuple[Callable[[], bool], Callable[[], AWSCredentialsResolver]]
-_DEFAULT_SOURCES: Sequence[CredentialSource] = (
-    (_env_creds_available, _build_env_creds),
+_DEFAULT_SOURCES: Sequence[CredentialsSource] = (
+    EnvironmentCredentialsSource(),
+    IMDSCredentialsSource(),
 )
 
 
@@ -30,8 +27,14 @@ class CredentialsResolverChain(
 ):
     """Resolves AWS Credentials from system environment variables."""
 
-    def __init__(self, *, sources: Sequence[CredentialSource] = _DEFAULT_SOURCES):
-        self._sources: Sequence[CredentialSource] = sources
+    def __init__(
+        self,
+        *,
+        config: AwsCredentialsConfig,
+        sources: Sequence[CredentialsSource] = _DEFAULT_SOURCES,
+    ):
+        self._config = config
+        self._sources: Sequence[CredentialsSource] = sources
         self._credentials_resolver: AWSCredentialsResolver | None = None
 
     async def get_identity(
@@ -43,8 +46,8 @@ async def get_identity(
             )
 
         for source in self._sources:
-            if source[0]():
-                self._credentials_resolver = source[1]()
+            if source.is_available(config=self._config):
+                self._credentials_resolver = source.build_resolver(config=self._config)
                 return await self._credentials_resolver.get_identity(
                     identity_properties=identity_properties
                 )

packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/environment.py

@@ -6,7 +6,12 @@
 from smithy_core.exceptions import SmithyIdentityException
 from smithy_core.interfaces.identity import IdentityProperties
 
-from ..identity import AWSCredentialsIdentity
+from smithy_aws_core.credentials_resolvers.interfaces import (
+    AwsCredentialsConfig,
+    CredentialsSource,
+)
+
+from ..identity import AWSCredentialsIdentity, AWSCredentialsResolver
 
 
 class EnvironmentCredentialsResolver(
@@ -41,3 +46,13 @@ async def get_identity(
         )
 
         return self._credentials
+
+
+class EnvironmentCredentialsSource(CredentialsSource):
+    def is_available(self, config: AwsCredentialsConfig) -> bool:
+        return (
+            "AWS_ACCESS_KEY_ID" in os.environ and "AWS_SECRET_ACCESS_KEY" in os.environ
+        )
+
+    def build_resolver(self, config: AwsCredentialsConfig) -> AWSCredentialsResolver:
+        return EnvironmentCredentialsResolver()

packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/imds.py

@@ -17,8 +17,13 @@
 from smithy_http.aio import HTTPRequest
 from smithy_http.aio.interfaces import HTTPClient
 
+from smithy_aws_core.credentials_resolvers.interfaces import (
+    AwsCredentialsConfig,
+    CredentialsSource,
+)
+
 from .. import __version__
-from ..identity import AWSCredentialsIdentity
+from ..identity import AWSCredentialsIdentity, AWSCredentialsResolver
 
 _USER_AGENT_FIELD = Field(
     name="User-Agent",
@@ -235,3 +240,13 @@ async def get_identity(
             account_id=account_id,
         )
         return self._credentials
+
+
+class IMDSCredentialsSource(CredentialsSource):
+    def is_available(self, config: AwsCredentialsConfig) -> bool:
+        # IMDS credentials should always be the last in the chain
+        # We cannot check if they available without actually making a call
+        return True
+
+    def build_resolver(self, config: AwsCredentialsConfig) -> AWSCredentialsResolver:
+        return IMDSCredentialsResolver(http_client=config.http_client)

packages/smithy-aws-core/src/smithy_aws_core/credentials_resolvers/interfaces.py

@@ -0,0 +1,24 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+from typing import Protocol
+
+from smithy_http.aio.interfaces import HTTPClient
+
+from smithy_aws_core.identity import AWSCredentialsResolver
+
+
+class AwsCredentialsConfig(Protocol):
+    """Configuration required for resolving credentials."""
+
+    http_client: HTTPClient
+    """A static endpoint to use for the request."""
+
+
+class CredentialsSource(Protocol):
+    def is_available(self, config: AwsCredentialsConfig) -> bool:
+        """Returns True if credentials are available from this source."""
+        ...
+
+    def build_resolver(self, config: AwsCredentialsConfig) -> AWSCredentialsResolver:
+        """Builds a credentials resolver for the given configuration."""
+        ...

packages/smithy-aws-core/tests/unit/credentials_resolvers/test_credentials_resolver_chain.py

@@ -1,23 +1,44 @@
+from dataclasses import dataclass
+from unittest.mock import Mock
+
 import pytest
 from smithy_aws_core.credentials_resolvers import (
     CredentialsResolverChain,
     StaticCredentialsResolver,
 )
-from smithy_aws_core.identity import AWSCredentialsIdentity
+from smithy_aws_core.credentials_resolvers.environment import (
+    EnvironmentCredentialsSource,
+)
+from smithy_aws_core.credentials_resolvers.interfaces import (
+    AwsCredentialsConfig,
+    CredentialsSource,
+)
+from smithy_aws_core.identity import AWSCredentialsIdentity, AWSCredentialsResolver
 from smithy_core.exceptions import SmithyIdentityException
 from smithy_core.interfaces.identity import IdentityProperties
+from smithy_http.aio.interfaces import HTTPClient
+
+
+@dataclass
+class Config:
+    http_client: HTTPClient
+
+    def __init__(self):
+        self.http_client = Mock(spec=HTTPClient)  # type: ignore
 
 
 async def test_no_sources_resolve():
-    resolver_chain = CredentialsResolverChain(sources=[])
+    resolver_chain = CredentialsResolverChain(sources=[], config=Config())
     with pytest.raises(SmithyIdentityException):
         await resolver_chain.get_identity(identity_properties=IdentityProperties())
 
 
 async def test_env_credentials_resolver_not_set(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.delenv("AWS_ACCESS_KEY_ID", raising=False)
     monkeypatch.delenv("AWS_SECRET_ACCESS_KEY", raising=False)
-    resolver_chain = CredentialsResolverChain()
+    resolver_chain = CredentialsResolverChain(
+        sources=[EnvironmentCredentialsSource()], config=Config()
+    )
 
     with pytest.raises(SmithyIdentityException):
         await resolver_chain.get_identity(identity_properties=IdentityProperties())
@@ -26,7 +47,9 @@ async def test_env_credentials_resolver_not_set(monkeypatch: pytest.MonkeyPatch)
 async def test_env_credentials_resolver_partial(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setenv("AWS_ACCESS_KEY_ID", "akid")
     monkeypatch.delenv("AWS_SECRET_ACCESS_KEY", raising=False)
-    resolver_chain = CredentialsResolverChain()
+    resolver_chain = CredentialsResolverChain(
+        sources=[EnvironmentCredentialsSource()], config=Config()
+    )
 
     with pytest.raises(SmithyIdentityException):
         await resolver_chain.get_identity(identity_properties=IdentityProperties())
@@ -35,7 +58,9 @@ async def test_env_credentials_resolver_partial(monkeypatch: pytest.MonkeyPatch)
 async def test_env_credentials_resolver_success(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setenv("AWS_ACCESS_KEY_ID", "akid")
     monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret")
-    resolver_chain = CredentialsResolverChain()
+    resolver_chain = CredentialsResolverChain(
+        sources=[EnvironmentCredentialsSource()], config=Config()
+    )
 
     credentials = await resolver_chain.get_identity(
         identity_properties=IdentityProperties()
@@ -50,8 +75,19 @@ async def test_custom_sources_with_static_credentials():
         secret_access_key="static_secret",
     )
     static_resolver = StaticCredentialsResolver(credentials=static_credentials)
+
+    class TestStaticSource(CredentialsSource):
+        def is_available(self, config: AwsCredentialsConfig) -> bool:
+            return True
+
+        def build_resolver(
+            self, config: AwsCredentialsConfig
+        ) -> AWSCredentialsResolver:
+            return static_resolver
+
     resolver_chain = CredentialsResolverChain(
-        sources=[(lambda: False, lambda: None), (lambda: True, lambda: static_resolver)]  # type: ignore
+        sources=[TestStaticSource()],
+        config=Config(),  # type: ignore
     )
 
     credentials = await resolver_chain.get_identity(