(airflow) branch main updated: Fix the way to get STS endpoint in EKS hook (#45520)

[vincbeck]

This is an automated email from the ASF dual-hosted git repository.

vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 103df61bde6 Fix the way to get STS endpoint in EKS hook (#45520)
103df61bde6 is described below

commit 103df61bde6e98de2466f42e76b4ac3bcc4ab8b5
Author: Vincent <97131062+vincb...@users.noreply.github.com>
AuthorDate: Thu Jan 9 13:20:12 2025 -0500

    Fix the way to get STS endpoint in EKS hook (#45520)
---
 providers/src/airflow/providers/amazon/aws/hooks/eks.py | 4 ++--
 providers/tests/amazon/aws/hooks/test_eks.py            | 8 +++-----
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/providers/src/airflow/providers/amazon/aws/hooks/eks.py 
b/providers/src/airflow/providers/amazon/aws/hooks/eks.py
index d48c103505a..4e8c0ca7ad6 100644
--- a/providers/src/airflow/providers/amazon/aws/hooks/eks.py
+++ b/providers/src/airflow/providers/amazon/aws/hooks/eks.py
@@ -32,6 +32,7 @@ from botocore.exceptions import ClientError
 from botocore.signers import RequestSigner
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.hooks.sts import StsHook
 from airflow.utils import yaml
 from airflow.utils.json import AirflowJsonEncoder
 
@@ -612,8 +613,7 @@ class EksHook(AwsBaseHook):
     def fetch_access_token_for_cluster(self, eks_cluster_name: str) -> str:
         session = self.get_session()
         service_id = self.conn.meta.service_model.service_id
-        sts_client = session.client("sts")
-        sts_url = 
f"{sts_client.meta.endpoint_url}/?Action=GetCallerIdentity&Version=2011-06-15"
+        sts_url = 
f"{StsHook().conn_client_meta.endpoint_url}/?Action=GetCallerIdentity&Version=2011-06-15"
 
         signer = RequestSigner(
             service_id=service_id,
diff --git a/providers/tests/amazon/aws/hooks/test_eks.py 
b/providers/tests/amazon/aws/hooks/test_eks.py
index 06cc7ddab53..10a93790ac6 100644
--- a/providers/tests/amazon/aws/hooks/test_eks.py
+++ b/providers/tests/amazon/aws/hooks/test_eks.py
@@ -22,7 +22,6 @@ from copy import deepcopy
 from pathlib import Path
 from typing import TYPE_CHECKING
 from unittest import mock
-from unittest.mock import Mock
 from urllib.parse import urlsplit
 
 import pytest
@@ -1284,14 +1283,13 @@ class TestEksHook:
             }
 
     @mock.patch("airflow.providers.amazon.aws.hooks.eks.RequestSigner")
+    @mock.patch("airflow.providers.amazon.aws.hooks.eks.StsHook")
     @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.conn")
     
@mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.get_session")
-    def test_fetch_access_token_for_cluster(self, mock_get_session, mock_conn, 
mock_signer):
+    def test_fetch_access_token_for_cluster(self, mock_get_session, mock_conn, 
mock_sts_hook, mock_signer):
         mock_signer.return_value.generate_presigned_url.return_value = 
"http://example.com";
         mock_get_session.return_value.region_name = "us-east-1"
-        client = Mock()
-        client.meta.endpoint_url = "https://sts.us-east-1.amazonaws.com";
-        mock_get_session.return_value.client.return_value = client
+        mock_sts_hook.return_value.conn_client_meta.endpoint_url = 
"https://sts.us-east-1.amazonaws.com";
         hook = EksHook()
         token = 
hook.fetch_access_token_for_cluster(eks_cluster_name="test-cluster")
         mock_signer.assert_called_once_with(