unit tests

tests/functional/test_credentials.py

@@ -790,7 +790,9 @@ def assert_session_credentials(self, expected_params, **kwargs):
         expected_creds = self.create_random_credentials()
         response = self.create_assume_role_response(expected_creds)
         session = StubbedSession(**kwargs)
-        stubber = session.stub('sts')
+        stubber = session.stub(
+            'sts', config=Config(signature_version=UNSIGNED)
+        )
         stubber.add_response(
             'assume_role_with_web_identity', response, expected_params
         )

tests/unit/data/endpoints/valid-rules/aws-account-id.json

@@ -0,0 +1,83 @@
+{
+  "parameters": {
+    "Region": {
+      "type": "string",
+      "builtIn": "AWS::Region",
+      "documentation": "The region to dispatch this request, eg. `us-east-1`."
+    },
+    "AccountId": {
+      "type": "string",
+      "builtIn": "AWS::Auth::AccountId",
+      "documentation": "The account ID to dispatch this request, eg. `us-east-1`."
+    }
+  },
+  "rules": [
+    {
+      "documentation": "Template the account ID into the URI when account ID is set",
+      "conditions": [
+        {
+          "fn": "isSet",
+          "argv": [
+            {
+              "ref": "AccountId"
+            }
+          ]
+        },
+        {
+          "fn": "isSet",
+          "argv": [
+            {
+              "ref": "Region"
+            }
+          ]
+        }
+      ],
+      "endpoint": {
+        "url": "https://{AccountId}.amazonaws.com",
+        "properties": {
+          "authSchemes": [
+            {
+              "name": "sigv4",
+              "signingName": "serviceName",
+              "signingRegion": "{Region}"
+            }
+          ]
+        }
+      },
+      "type": "endpoint"
+    },
+    {
+      "documentation": "Fallback when account ID isn't set",
+      "conditions": [
+        {
+          "fn": "isSet",
+          "argv": [
+            {
+              "ref": "Region"
+            }
+          ]
+        }
+      ],
+      "endpoint": {
+        "url": "https://amazonaws.com",
+        "properties": {
+          "authSchemes": [
+            {
+              "name": "sigv4",
+              "signingName": "serviceName",
+              "signingRegion": "{Region}"
+            }
+          ]
+        }
+      },
+      "type": "endpoint"
+    },
+    {
+      "documentation": "fallback when region is unset",
+      "conditions": [],
+      "error": "Region must be set to resolve a valid endpoint",
+      "type": "error"
+    }
+  ],
+  "version": "1.3"
+}

tests/unit/test_args.py

@@ -679,6 +679,7 @@ def test_builtins_defaults(self):
             bins['AWS::S3::DisableMultiRegionAccessPoints'], False
         )
         self.assertEqual(bins['SDK::Endpoint'], None)
+        self.assertEqual(bins['AWS::Auth::AccountId'], None)
 
     def test_aws_region(self):
         bins = self.call_compute_endpoint_resolver_builtin_defaults(

tests/unit/test_endpoint_provider.py

@@ -18,6 +18,9 @@
 
 import pytest
 
+from botocore import UNSIGNED
+from botocore.config import Config
+from botocore.credentials import Credentials
 from botocore.endpoint_provider import (
     EndpointProvider,
     EndpointRule,
@@ -28,12 +31,14 @@
     TreeRule,
 )
 from botocore.exceptions import (
+    AccountIDNotFound,
     EndpointResolutionError,
+    InvalidConfigError,
     MissingDependencyException,
     UnknownSignatureVersionError,
 )
 from botocore.loaders import Loader
-from botocore.regions import EndpointRulesetResolver
+from botocore.regions import EndpointResolverBuiltins, EndpointRulesetResolver
 from tests import requires_crt
 
 REGION_TEMPLATE = "{Region}"
@@ -511,3 +516,156 @@ def test_aws_is_virtual_hostable_s3_bucket_allow_subdomains(
         rule_lib.aws_is_virtual_hostable_s3_bucket(bucket, True)
         == expected_value
     )
+
+
+def account_id_ruleset():
+    rule_path = os.path.join(
+        os.path.dirname(__file__),
+        "data",
+        "endpoints",
+        "valid-rules",
+        "aws-account-id.json",
+    )
+    with open(rule_path) as f:
+        return json.load(f)
+
+
+def _service_model_empty_context_params():
+    service_model = Mock()
+    service_model.client_context_parameters = []
+    return service_model
+
+
+@pytest.fixture
+def operation_model_empty_context_params():
+    operation_model = Mock()
+    operation_model.static_context_parameters = []
+    operation_model.context_parameters = []
+    return operation_model
+
+
+ACCOUNT_ID_RULESET = account_id_ruleset()
+SERVICE_MODEL_EMPTY_CONTEXT_PARAMS = _service_model_empty_context_params()
+BUILTINS_WITH_ACCOUNT_ID = {
+    EndpointResolverBuiltins.AWS_REGION: "us-west-2",
+    EndpointResolverBuiltins.AWS_ACCOUNT_ID: None,
+}
+STATIC_CREDENTIALS = Credentials(
+    access_key="access_key",
+    secret_key="secret_key",
+    token="token",
+    account_id="1234567890",
+)
+
+
+def resolver_with_account_id_builtin(credentials, auth_scheme):
+    return EndpointRulesetResolver(
+        endpoint_ruleset_data=ACCOUNT_ID_RULESET,
+        partition_data={},
+        service_model=SERVICE_MODEL_EMPTY_CONTEXT_PARAMS,
+        builtins=BUILTINS_WITH_ACCOUNT_ID,
+        client_context=None,
+        event_emitter=Mock(),
+        use_ssl=True,
+        credentials=credentials,
+        requested_auth_scheme=auth_scheme,
+    )
+
+
+ACT_ID_REQUIRED_CONTEXT = {
+    'client_config': Config(account_id_endpoint_mode="required")
+}
+ACT_ID_PREFERRED_CONTEXT = {
+    'client_config': Config(account_id_endpoint_mode="preferred")
+}
+
+URL_NO_ACCOUNT_ID = "https://amazonaws.com"
+URL_WITH_ACCOUNT_ID = "https://1234567890.amazonaws.com"
+
+
+@pytest.mark.parametrize(
+    "credentials, auth_scheme, request_context, expected_url",
+    [
+        (
+            STATIC_CREDENTIALS,
+            None,
+            ACT_ID_REQUIRED_CONTEXT,
+            URL_WITH_ACCOUNT_ID,
+        ),
+        (
+            STATIC_CREDENTIALS,
+            None,
+            ACT_ID_PREFERRED_CONTEXT,
+            URL_WITH_ACCOUNT_ID,
+        ),
+        (
+            STATIC_CREDENTIALS,
+            None,
+            {'client_config': Config(account_id_endpoint_mode="disabled")},
+            URL_NO_ACCOUNT_ID,
+        ),
+        (
+            STATIC_CREDENTIALS,
+            UNSIGNED,
+            ACT_ID_REQUIRED_CONTEXT,
+            URL_NO_ACCOUNT_ID,
+        ),
+        (
+            STATIC_CREDENTIALS,
+            None,
+            {**ACT_ID_REQUIRED_CONTEXT, "is_presign_request": True},
+            URL_NO_ACCOUNT_ID,
+        ),
+        (
+            None,
+            None,
+            ACT_ID_PREFERRED_CONTEXT,
+            URL_NO_ACCOUNT_ID,
+        ),
+        (
+            Credentials(access_key="foo", secret_key="bar", token="baz"),
+            None,
+            ACT_ID_PREFERRED_CONTEXT,
+            URL_NO_ACCOUNT_ID,
+        ),
+    ],
+)
+def test_account_id_builtin(
+    operation_model_empty_context_params,
+    credentials,
+    auth_scheme,
+    request_context,
+    expected_url,
+):
+    resolver = resolver_with_account_id_builtin(credentials, auth_scheme)
+    endpoint = resolver.construct_endpoint(
+        operation_model=operation_model_empty_context_params,
+        request_context=request_context,
+        call_args={},
+    )
+    assert endpoint.url == expected_url
+
+
+def test_invalid_account_id_endpoint_mode_raises(
+    operation_model_empty_context_params,
+):
+    resolver = resolver_with_account_id_builtin(STATIC_CREDENTIALS, None)
+    with pytest.raises(InvalidConfigError):
+        resolver.construct_endpoint(
+            operation_model=operation_model_empty_context_params,
+            request_context={
+                'client_config': Config(account_id_endpoint_mode="foo")
+            },
+            call_args={},
+        )
+
+
+def test_account_id_not_found_raises(operation_model_empty_context_params):
+    credentials = Credentials(access_key="foo", secret_key="bar", token="baz")
+    resolver = resolver_with_account_id_builtin(credentials, None)
+    with pytest.raises(AccountIDNotFound):
+        resolver.construct_endpoint(
+            operation_model=operation_model_empty_context_params,
+            request_context=ACT_ID_REQUIRED_CONTEXT,
+            call_args={},
+        )