From 2d1b2b26c6ba11732d294cfe9adf8c79233da5bc Mon Sep 17 00:00:00 2001
From: Andrey Anshin <Andrey.Anshin@taragol.is>
Date: Sun, 24 Sep 2023 13:48:01 +0400
Subject: [PATCH] Allow setup `endpoint_url` per-service in AWS Connection

---
 .../providers/amazon/aws/hooks/base_aws.py    |  70 ++++++---
 .../amazon/aws/utils/connection_wrapper.py    |  42 ++++-
 .../connections/aws.rst                       |  36 ++++-
 .../amazon/aws/hooks/test_base_aws.py         | 145 ++++++++++++------
 .../aws/utils/test_connection_wrapper.py      |  77 +++++++++-
 5 files changed, 299 insertions(+), 71 deletions(-)

diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py
index 02214f406db56..742ec5293237d 100644
--- a/airflow/providers/amazon/aws/hooks/base_aws.py
+++ b/airflow/providers/amazon/aws/hooks/base_aws.py
@@ -246,7 +246,11 @@ def _refresh_credentials(self) -> dict[str, Any]:
         if assume_role_method not in ("assume_role", "assume_role_with_saml"):
             raise NotImplementedError(f"assume_role_method={assume_role_method} not expected")
 
-        sts_client = self.basic_session.client("sts", config=self.config)
+        sts_client = self.basic_session.client(
+            "sts",
+            config=self.config,
+            endpoint_url=self.conn.get_service_endpoint_url("sts", sts_connection_assume=True),
+        )
 
         if assume_role_method == "assume_role":
             sts_response = self._assume_role(sts_client=sts_client)
@@ -558,10 +562,33 @@ def conn_config(self) -> AwsConnectionWrapper:
             conn=connection, region_name=self._region_name, botocore_config=self._config, verify=self._verify
         )
 
+    def _resolve_service_name(self, is_resource_type: bool = False) -> str:
+        """Resolve service name based on type or raise an error."""
+        if exactly_one(self.client_type, self.resource_type):
+            # It is possible to write simple conditions, however it make mypy unhappy.
+            if self.client_type:
+                if is_resource_type:
+                    raise LookupError("Requested `resource_type`, but `client_type` was set instead.")
+                return self.client_type
+            elif self.resource_type:
+                if not is_resource_type:
+                    raise LookupError("Requested `client_type`, but `resource_type` was set instead.")
+                return self.resource_type
+
+        raise ValueError(
+            f"Either client_type={self.client_type!r} or "
+            f"resource_type={self.resource_type!r} must be provided, not both."
+        )
+
+    @property
+    def service_name(self) -> str:
+        """Extracted botocore/boto3 service name from hook parameters."""
+        return self._resolve_service_name(is_resource_type=bool(self.resource_type))
+
     @property
     def service_config(self) -> dict:
-        service_name = self.client_type or self.resource_type
-        return self.conn_config.get_service_config(service_name)
+        """Config for hook-specific service from AWS Connection."""
+        return self.conn_config.get_service_config(service_name=self.service_name)
 
     @property
     def region_name(self) -> str | None:
@@ -609,19 +636,20 @@ def get_client_type(
         deferrable: bool = False,
     ) -> boto3.client:
         """Get the underlying boto3 client using boto3 session."""
-        client_type = self.client_type
+        service_name = self._resolve_service_name(is_resource_type=False)
         session = self.get_session(region_name=region_name, deferrable=deferrable)
+        endpoint_url = self.conn_config.get_service_endpoint_url(service_name=service_name)
         if not isinstance(session, boto3.session.Session):
             return session.create_client(
-                client_type,
-                endpoint_url=self.conn_config.endpoint_url,
+                service_name=service_name,
+                endpoint_url=endpoint_url,
                 config=self._get_config(config),
                 verify=self.verify,
             )
 
         return session.client(
-            client_type,
-            endpoint_url=self.conn_config.endpoint_url,
+            service_name=service_name,
+            endpoint_url=endpoint_url,
             config=self._get_config(config),
             verify=self.verify,
         )
@@ -632,11 +660,11 @@ def get_resource_type(
         config: Config | None = None,
     ) -> boto3.resource:
         """Get the underlying boto3 resource using boto3 session."""
-        resource_type = self.resource_type
+        service_name = self._resolve_service_name(is_resource_type=True)
         session = self.get_session(region_name=region_name)
         return session.resource(
-            resource_type,
-            endpoint_url=self.conn_config.endpoint_url,
+            service_name=service_name,
+            endpoint_url=self.conn_config.get_service_endpoint_url(service_name=service_name),
             config=self._get_config(config),
             verify=self.verify,
         )
@@ -648,15 +676,9 @@ def conn(self) -> BaseAwsConnection:
 
         :return: boto3.client or boto3.resource
         """
-        if not exactly_one(self.client_type, self.resource_type):
-            raise ValueError(
-                f"Either client_type={self.client_type!r} or "
-                f"resource_type={self.resource_type!r} must be provided, not both."
-            )
-        elif self.client_type:
+        if self.client_type:
             return self.get_client_type(region_name=self.region_name)
-        else:
-            return self.get_resource_type(region_name=self.region_name)
+        return self.get_resource_type(region_name=self.region_name)
 
     @property
     def async_conn(self):
@@ -730,7 +752,10 @@ def expand_role(self, role: str, region_name: str | None = None) -> str:
         else:
             session = self.get_session(region_name=region_name)
             _client = session.client(
-                "iam", endpoint_url=self.conn_config.endpoint_url, config=self.config, verify=self.verify
+                service_name="iam",
+                endpoint_url=self.conn_config.get_service_endpoint_url("iam"),
+                config=self.config,
+                verify=self.verify,
             )
             return _client.get_role(RoleName=role)["Role"]["Arn"]
 
@@ -799,10 +824,9 @@ def test_connection(self):
         """
         try:
             session = self.get_session()
-            test_endpoint_url = self.conn_config.extra_config.get("test_endpoint_url")
             conn_info = session.client(
-                "sts",
-                endpoint_url=test_endpoint_url,
+                service_name="sts",
+                endpoint_url=self.conn_config.get_service_endpoint_url("sts", sts_test_connection=True),
             ).get_caller_identity()
             metadata = conn_info.pop("ResponseMetadata", {})
             if metadata.get("HTTPStatusCode") != 200:
diff --git a/airflow/providers/amazon/aws/utils/connection_wrapper.py b/airflow/providers/amazon/aws/utils/connection_wrapper.py
index 1520c3e1fc06f..f0b8d7c1c26b4 100644
--- a/airflow/providers/amazon/aws/utils/connection_wrapper.py
+++ b/airflow/providers/amazon/aws/utils/connection_wrapper.py
@@ -122,12 +122,48 @@ class AwsConnectionWrapper(LoggingMixin):
     assume_role_method: str | None = field(init=False, default=None)
     assume_role_kwargs: dict[str, Any] = field(init=False, default_factory=dict)
 
+    # Per AWS Service configuration dictionary where key is name of boto3 ``service_name``
+    service_config: dict[str, dict[str, Any]] = field(init=False, default_factory=dict)
+
     @cached_property
     def conn_repr(self):
         return f"AWS Connection (conn_id={self.conn_id!r}, conn_type={self.conn_type!r})"
 
-    def get_service_config(self, service_name):
-        return self.extra_dejson.get("service_config", {}).get(service_name, {})
+    def get_service_config(self, service_name: str) -> dict[str, Any]:
+        """Get AWS Service related config dictionary.
+
+        :param service_name: Name of botocore/boto3 service.
+        """
+        return self.service_config.get(service_name, {})
+
+    def get_service_endpoint_url(
+        self, service_name: str, *, sts_connection_assume: bool = False, sts_test_connection: bool = False
+    ) -> str | None:
+        service_config = self.get_service_config(service_name=service_name)
+        global_endpoint_url = self.endpoint_url
+
+        if service_name == "sts" and True in (sts_connection_assume, sts_test_connection):
+            # There are different logics exists historically for STS Client
+            # 1. For assume role we never use global endpoint_url
+            # 2. For test connection we also use undocumented `test_endpoint`\
+            # 3. For STS as service we might use endpoint_url (default for other services)
+            global_endpoint_url = None
+            if sts_connection_assume and sts_test_connection:
+                raise AirflowException(
+                    "Can't resolve STS endpoint when both "
+                    "`sts_connection` and `sts_test_connection` set to True."
+                )
+            elif sts_test_connection:
+                if "test_endpoint_url" in self.extra_config:
+                    warnings.warn(
+                        "extra['test_endpoint_url'] is deprecated and will be removed in a future release."
+                        " Please set `endpoint_url` in `service_config.sts` within `extras`.",
+                        AirflowProviderDeprecationWarning,
+                        stacklevel=2,
+                    )
+                    global_endpoint_url = self.extra_config["test_endpoint_url"]
+
+        return service_config.get("endpoint_url", global_endpoint_url)
 
     def __post_init__(self, conn: Connection):
         if isinstance(conn, type(self)):
@@ -182,6 +218,8 @@ def __post_init__(self, conn: Connection):
             )
 
         extra = deepcopy(conn.extra_dejson)
+        self.service_config = extra.get("service_config", {})
+
         session_kwargs = extra.get("session_kwargs", {})
         if session_kwargs:
             warnings.warn(
diff --git a/docs/apache-airflow-providers-amazon/connections/aws.rst b/docs/apache-airflow-providers-amazon/connections/aws.rst
index c12e24e76541e..0b87e7be93fd8 100644
--- a/docs/apache-airflow-providers-amazon/connections/aws.rst
+++ b/docs/apache-airflow-providers-amazon/connections/aws.rst
@@ -122,7 +122,9 @@ Extra (optional)
     * ``config_kwargs``: Additional **kwargs** used to construct a
       `botocore.config.Config <https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html>`__.
       To anonymously access public AWS resources (equivalent of `signature_version=botocore.UNSGINED`), set `"signature_version"="unsigned"` within `config_kwargs`.
-    * ``endpoint_url``: Endpoint URL for the connection.
+    * ``endpoint_url``: Global Endpoint URL for the connection. You could specify endpoint url per AWS service by utilize
+      ``service_config``, for more details please refer to :ref:`howto/connection:aws:per-service-endpoint-configuration`
+
     * ``verify``: Whether or not to verify SSL certificates.
 
     The following extra parameters used for specific AWS services:
@@ -343,6 +345,38 @@ The following settings may be used within the ``assume_role_with_saml`` containe
 Per-service configuration
 ^^^^^^^^^^^^^^^^^^^^^^^^^
 
+.. _howto/connection:aws:per-service-endpoint-configuration:
+
+AWS Service Endpoint URL configuration
+""""""""""""""""""""""""""""""""""""""
+
+To use ``endpoint_url`` per specific AWS service in the single connection you might setup it in service config.
+For enforce to default ``botocore``/``boto3`` behaviour you might set value to ``null``.
+The precedence rules are as follows:
+
+  1. ``endpoint_url`` specified per service level.
+  2. ``endpoint_url`` specified in root level of connection extra. Please note that **sts** client which are
+     uses in assume role or test connection do not use global parameter.
+  3. Default ``botocore``/``boto3`` behaviour
+
+
+.. code-block:: json
+
+    {
+      "endpoint_url": "s3.amazonaws.com"
+      "service_config": {
+        "s3": {
+          "endpoint_url": "https://s3.eu-west-1.amazonaws.com"
+        },
+        "sts": {
+          "endpoint_url": "https://sts.eu-west-2.amazonaws.com"
+        },
+        "ec2": {
+          "endpoint_url": null
+        }
+      }
+    }
+
 S3 Bucket configurations
 """"""""""""""""""""""""
 
diff --git a/tests/providers/amazon/aws/hooks/test_base_aws.py b/tests/providers/amazon/aws/hooks/test_base_aws.py
index 660190a6e908f..57350cb74115b 100644
--- a/tests/providers/amazon/aws/hooks/test_base_aws.py
+++ b/tests/providers/amazon/aws/hooks/test_base_aws.py
@@ -21,6 +21,7 @@
 import json
 import os
 from base64 import b64encode
+from contextlib import nullcontext
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
 from unittest import mock
@@ -38,7 +39,7 @@
 from moto import mock_dynamodb, mock_emr, mock_iam, mock_sts
 from moto.core import DEFAULT_ACCOUNT_ID
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.models.connection import Connection
 from airflow.providers.amazon.aws.hooks.base_aws import (
     AwsBaseHook,
@@ -440,22 +441,24 @@ def test_user_agent_dag_run_key_is_hashed_correctly(self, _, env_var, expected_v
 
             assert UUID(dag_run_key).version == expected_version
 
+    @pytest.mark.parametrize(
+        "sts_endpoint",
+        [
+            pytest.param(None, id="not-set"),
+            pytest.param("https://foo.bar/spam/egg", id="custom"),
+        ],
+    )
     @mock.patch.object(AwsBaseHook, "get_connection")
     @mock_sts
-    def test_assume_role(self, mock_get_connection):
+    def test_assume_role(self, mock_get_connection, sts_endpoint):
         aws_conn_id = "aws/test"
         role_arn = "arn:aws:iam::123456:role/role_arn"
         slugified_role_session_name = "airflow_aws-test"
 
-        mock_connection = Connection(
-            conn_id=aws_conn_id,
-            extra=json.dumps(
-                {
-                    "role_arn": role_arn,
-                }
-            ),
-        )
-        mock_get_connection.return_value = mock_connection
+        fake_conn_extra = {"role_arn": role_arn, "endpoint_url": "https://example.org"}
+        if sts_endpoint:
+            fake_conn_extra["service_config"] = {"sts": {"endpoint_url": sts_endpoint}}
+        mock_get_connection.return_value = Connection(conn_id=aws_conn_id, extra=fake_conn_extra)
 
         def mock_assume_role(**kwargs):
             assert kwargs["RoleArn"] == role_arn
@@ -490,7 +493,7 @@ def mock_assume_role(**kwargs):
             hook.get_client_type("s3")
 
         calls_assume_role = [
-            mock.call.session.Session().client("sts", config=mock.ANY),
+            mock.call.session.Session().client("sts", config=mock.ANY, endpoint_url=sts_endpoint),
             mock.call.session.Session()
             .client()
             .assume_role(
@@ -610,34 +613,39 @@ def test_get_credentials_from_token_file(
         assert mock_creds_fetcher_kwargs["web_identity_token_loader"]() == "TOKEN"
         assert mock_open_.call_args.args[0] == "/my-token-path"
 
+    @pytest.mark.parametrize(
+        "sts_endpoint",
+        [
+            pytest.param(None, id="not-set"),
+            pytest.param("https://foo.bar/spam/egg", id="custom"),
+        ],
+    )
     @mock.patch.object(AwsBaseHook, "get_connection")
     @mock_sts
-    def test_assume_role_with_saml(self, mock_get_connection):
+    def test_assume_role_with_saml(self, mock_get_connection, sts_endpoint):
         idp_url = "https://my-idp.local.corp"
         principal_arn = "principal_arn_1234567890"
         role_arn = "arn:aws:iam::123456:role/role_arn"
         xpath = "1234"
         duration_seconds = 901
 
-        mock_connection = Connection(
-            conn_id=MOCK_AWS_CONN_ID,
-            extra=json.dumps(
-                {
-                    "role_arn": role_arn,
-                    "assume_role_method": "assume_role_with_saml",
-                    "assume_role_with_saml": {
-                        "principal_arn": principal_arn,
-                        "idp_url": idp_url,
-                        "idp_auth_method": "http_spegno_auth",
-                        "mutual_authentication": "REQUIRED",
-                        "saml_response_xpath": xpath,
-                        "log_idp_response": True,
-                    },
-                    "assume_role_kwargs": {"DurationSeconds": duration_seconds},
-                }
-            ),
-        )
-        mock_get_connection.return_value = mock_connection
+        fake_conn_extra = {
+            "role_arn": role_arn,
+            "assume_role_method": "assume_role_with_saml",
+            "assume_role_with_saml": {
+                "principal_arn": principal_arn,
+                "idp_url": idp_url,
+                "idp_auth_method": "http_spegno_auth",
+                "mutual_authentication": "REQUIRED",
+                "saml_response_xpath": xpath,
+                "log_idp_response": True,
+            },
+            "assume_role_kwargs": {"DurationSeconds": duration_seconds},
+            "endpoint_url": "https://example.org",
+        }
+        if sts_endpoint:
+            fake_conn_extra["service_config"] = {"sts": {"endpoint_url": sts_endpoint}}
+        mock_get_connection.return_value = Connection(conn_id=MOCK_AWS_CONN_ID, extra=fake_conn_extra)
 
         encoded_saml_assertion = b64encode(SAML_ASSERTION.encode("utf-8")).decode("utf-8")
 
@@ -693,7 +701,7 @@ def mock_assume_role_with_saml(**kwargs):
             mock_xpath.assert_called_once_with(xpath)
 
         calls_assume_role_with_saml = [
-            mock.call.session.Session().client("sts", config=mock.ANY),
+            mock.call.session.Session().client("sts", config=mock.ANY, endpoint_url=sts_endpoint),
             mock.call.session.Session()
             .client()
             .assume_role_with_saml(
@@ -822,6 +830,25 @@ def test_connection_aws_partition(self, conn_type, connection_uri, expected_part
 
             assert hook.conn_partition == expected_partition
 
+    @mock_dynamodb
+    def test_service_name(self):
+        client_hook = AwsBaseHook(aws_conn_id=None, client_type="dynamodb")
+        resource_hook = AwsBaseHook(aws_conn_id=None, resource_type="dynamodb")
+        # Should not raise any error here
+        invalid_hook = AwsBaseHook(aws_conn_id=None, client_type="dynamodb", resource_type="dynamodb")
+
+        assert client_hook.service_name == "dynamodb"
+        assert resource_hook.service_name == "dynamodb"
+
+        with pytest.raises(ValueError, match="Either client_type=.* or resource_type=.* must be provided"):
+            invalid_hook.service_name
+
+        with pytest.raises(LookupError, match="Requested `resource_type`, but `client_type` was set instead"):
+            client_hook._resolve_service_name(is_resource_type=True)
+
+        with pytest.raises(LookupError, match="Requested `client_type`, but `resource_type` was set instead"):
+            resource_hook._resolve_service_name(is_resource_type=False)
+
     @pytest.mark.parametrize(
         "client_type,resource_type",
         [
@@ -859,7 +886,7 @@ def test_hook_connection_test_failed(self, mock_boto3_session):
         result, message = hook.test_connection()
         assert not result
         assert message == json.dumps(response_metadata)
-        mock_sts_client.assert_called_once_with("sts", endpoint_url=None)
+        mock_sts_client.assert_called_once_with(service_name="sts", endpoint_url=None)
 
         def mock_error():
             raise ConnectionError("Test Error")
@@ -873,22 +900,52 @@ def mock_error():
 
         assert hook.client_type == "ec2"
 
-    @mock_sts
     @pytest.mark.parametrize(
-        "test_endpoint_url, result_url",
+        "sts_service_endpoint_url, test_endpoint_url, result_url",
         [
-            (None, "https://sts.amazonaws.com"),
-            ("https://sts.us-east-1.amazonaws.com", "https://sts.us-east-1.amazonaws.com"),
+            pytest.param(None, None, None, id="not-set"),
+            pytest.param(
+                "https://sts.service:1234", None, "https://sts.service:1234", id="sts-service-endpoint"
+            ),
+            pytest.param(
+                None, "http://deprecated.test", "http://deprecated.test", id="deprecated-test-parameter"
+            ),
+            pytest.param(
+                "https://sts.service:1234",
+                "http://deprecated.test",
+                "https://sts.service:1234",
+                id="mixin-resolve",
+            ),
         ],
     )
-    def test_hook_connection_endpoint_url_valid(self, test_endpoint_url, result_url):
+    @mock.patch("boto3.session.Session")
+    def test_hook_connection_endpoint_url_valid(
+        self, mock_boto3_session, sts_service_endpoint_url, test_endpoint_url, result_url, monkeypatch
+    ):
         """Test if test_endpoint_url is valid in test connection"""
-        conn = AwsConnectionWrapper.from_connection_metadata(conn_id=MOCK_AWS_CONN_ID)
-        sf = BaseSessionFactory(conn=conn)
-        session = sf.create_session()
-        client = session.client("sts", endpoint_url=test_endpoint_url)
 
-        assert client._endpoint.host == result_url
+        mock_sts_client = mock.MagicMock()
+        mock_boto3_session.return_value.client = mock_sts_client
+
+        warn_context = nullcontext()
+        fake_extra = {"endpoint_url": "https://test.conn:777/should/ignore/global/endpoint/url"}
+        if test_endpoint_url:
+            fake_extra["test_endpoint_url"] = test_endpoint_url
+            # If `test_endpoint_url` set than we raise warning message
+            warn_context = pytest.warns(
+                AirflowProviderDeprecationWarning, match=r"extra\['test_endpoint_url'\] is deprecated"
+            )
+        if sts_service_endpoint_url:
+            fake_extra["service_config"] = {"sts": {"endpoint_url": sts_service_endpoint_url}}
+
+        monkeypatch.setenv(
+            f"AIRFLOW_CONN_{MOCK_AWS_CONN_ID.upper()}", json.dumps({"conn_type": "aws", "extra": fake_extra})
+        )
+        hook = AwsBaseHook(aws_conn_id=MOCK_AWS_CONN_ID, client_type="eks")
+        with warn_context:
+            hook.test_connection()
+
+        mock_sts_client.assert_called_once_with(service_name="sts", endpoint_url=result_url)
 
     @mock.patch.dict(os.environ, {f"AIRFLOW_CONN_{MOCK_AWS_CONN_ID.upper()}": "aws://"})
     def test_conn_config_conn_id_exists(self):
diff --git a/tests/providers/amazon/aws/utils/test_connection_wrapper.py b/tests/providers/amazon/aws/utils/test_connection_wrapper.py
index e8f28e19b09f3..30fffb8fc0f45 100644
--- a/tests/providers/amazon/aws/utils/test_connection_wrapper.py
+++ b/tests/providers/amazon/aws/utils/test_connection_wrapper.py
@@ -23,7 +23,7 @@
 from botocore import UNSIGNED
 from botocore.config import Config
 
-from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.models import Connection
 from airflow.providers.amazon.aws.utils.connection_wrapper import AwsConnectionWrapper, _ConnectionMetadata
 
@@ -529,3 +529,78 @@ def test_get_wrapper_from_metadata(self, conn_id, profile_name, role_arn):
         assert wrap_conn.conn_id == conn_id
         assert wrap_conn.role_arn == role_arn
         assert wrap_conn.profile_name == profile_name
+
+    def test_get_service_config(self):
+        mock_conn = mock_connection_factory(
+            conn_id="foo-bar",
+            extra={
+                "service_config": {
+                    "sns": {"foo": "bar"},
+                    "s3": {"spam": "egg", "baz": "qux"},
+                    "dynamodb": None,
+                },
+            },
+        )
+        wrap_conn = AwsConnectionWrapper(conn=mock_conn)
+        assert wrap_conn.get_service_config("sns") == {"foo": "bar"}
+        assert wrap_conn.get_service_config("s3") == {"spam": "egg", "baz": "qux"}
+        assert wrap_conn.get_service_config("ec2") == {}
+        assert wrap_conn.get_service_config("dynamodb") is None
+
+    def test_get_service_endpoint_url(self):
+        fake_conn = mock_connection_factory(
+            conn_id="foo-bar",
+            extra={
+                "endpoint_url": "https://spam.egg",
+                "service_config": {
+                    "sns": {"endpoint_url": "https://foo.bar"},
+                    "ec2": {"endpoint_url": None},  # Enforce to boto3
+                },
+            },
+        )
+        wrap_conn = AwsConnectionWrapper(conn=fake_conn)
+        assert wrap_conn.get_service_endpoint_url("sns") == "https://foo.bar"
+        assert wrap_conn.get_service_endpoint_url("sts") == "https://spam.egg"
+        assert wrap_conn.get_service_endpoint_url("ec2") is None
+
+    @pytest.mark.parametrize(
+        "global_endpoint_url, sts_service_endpoint_url, expected_endpoint_url",
+        [
+            pytest.param(None, None, None, id="not-set"),
+            pytest.param("https://global.service", None, None, id="global-only"),
+            pytest.param(None, "https://sts.service:1234", "https://sts.service:1234", id="service-only"),
+            pytest.param(
+                "https://global.service", "https://sts.service:1234", "https://sts.service:1234", id="mixin"
+            ),
+        ],
+    )
+    def test_get_service_endpoint_url_sts(
+        self, global_endpoint_url, sts_service_endpoint_url, expected_endpoint_url
+    ):
+        fake_extra = {}
+        if global_endpoint_url:
+            fake_extra["endpoint_url"] = global_endpoint_url
+        if sts_service_endpoint_url:
+            fake_extra["service_config"] = {"sts": {"endpoint_url": sts_service_endpoint_url}}
+
+        fake_conn = mock_connection_factory(conn_id="foo-bar", extra=fake_extra)
+        wrap_conn = AwsConnectionWrapper(conn=fake_conn)
+        assert wrap_conn.get_service_endpoint_url("sts", sts_connection_assume=True) == expected_endpoint_url
+        assert wrap_conn.get_service_endpoint_url("sts", sts_test_connection=True) == expected_endpoint_url
+
+    def test_get_service_endpoint_url_sts_deprecated_test_connection(self):
+        fake_conn = mock_connection_factory(
+            conn_id="foo-bar",
+            extra={"endpoint_url": "https://spam.egg", "test_endpoint_url": "https://foo.bar"},
+        )
+        wrap_conn = AwsConnectionWrapper(conn=fake_conn)
+        warning_message = r"extra\['test_endpoint_url'\] is deprecated"
+        with pytest.warns(AirflowProviderDeprecationWarning, match=warning_message):
+            assert wrap_conn.get_service_endpoint_url("sts", sts_test_connection=True) == "https://foo.bar"
+
+    def test_get_service_endpoint_url_sts_unsupported(self):
+        wrap_conn = AwsConnectionWrapper(conn=mock_connection_factory())
+        with pytest.raises(AirflowException, match=r"Can't resolve STS endpoint when both"):
+            wrap_conn.get_service_endpoint_url("sts", sts_test_connection=True, sts_connection_assume=True)
+        # This check is only affects STS service endpoints
+        wrap_conn.get_service_endpoint_url("s3", sts_test_connection=True, sts_connection_assume=True)
