diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py index 72771a0c1f4b4..7f74447827da7 100644 --- a/airflow/providers/google/cloud/hooks/kubernetes_engine.py +++ b/airflow/providers/google/cloud/hooks/kubernetes_engine.py @@ -508,14 +508,15 @@ async def get_pod(self, name: str, namespace: str) -> V1Pod: :param name: Name of the pod. :param namespace: Name of the pod's namespace. """ - async with Token(scopes=self.scopes) as token: - async with self.get_conn(token) as connection: - v1_api = async_client.CoreV1Api(connection) - pod: V1Pod = await v1_api.read_namespaced_pod( - name=name, - namespace=namespace, - ) - return pod + async with self.service_file_as_context() as service_file: # type: ignore[attr-defined] + async with Token(scopes=self.scopes, service_file=service_file) as token: + async with self.get_conn(token) as connection: + v1_api = async_client.CoreV1Api(connection) + pod: V1Pod = await v1_api.read_namespaced_pod( + name=name, + namespace=namespace, + ) + return pod async def delete_pod(self, name: str, namespace: str): """Delete a pod. @@ -523,18 +524,21 @@ async def delete_pod(self, name: str, namespace: str): :param name: Name of the pod. :param namespace: Name of the pod's namespace. """ - async with Token(scopes=self.scopes) as token, self.get_conn(token) as connection: - try: - v1_api = async_client.CoreV1Api(connection) - await v1_api.delete_namespaced_pod( - name=name, - namespace=namespace, - body=client.V1DeleteOptions(), - ) - except async_client.ApiException as e: - # If the pod is already deleted - if e.status != 404: - raise + async with self.service_file_as_context() as service_file: # type: ignore[attr-defined] + async with Token(scopes=self.scopes, service_file=service_file) as token, self.get_conn( + token + ) as connection: + try: + v1_api = async_client.CoreV1Api(connection) + await v1_api.delete_namespaced_pod( + name=name, + namespace=namespace, + body=client.V1DeleteOptions(), + ) + except async_client.ApiException as e: + # If the pod is already deleted + if e.status != 404: + raise async def read_logs(self, name: str, namespace: str): """Read logs inside the pod while starting containers inside. @@ -547,19 +551,22 @@ async def read_logs(self, name: str, namespace: str): :param name: Name of the pod. :param namespace: Name of the pod's namespace. """ - async with Token(scopes=self.scopes) as token, self.get_conn(token) as connection: - try: - v1_api = async_client.CoreV1Api(connection) - logs = await v1_api.read_namespaced_pod_log( - name=name, - namespace=namespace, - follow=False, - timestamps=True, - ) - logs = logs.splitlines() - for line in logs: - self.log.info("Container logs from %s", line) - return logs - except HTTPError: - self.log.exception("There was an error reading the kubernetes API.") - raise + async with self.service_file_as_context() as service_file: # type: ignore[attr-defined] + async with Token(scopes=self.scopes, service_file=service_file) as token, self.get_conn( + token + ) as connection: + try: + v1_api = async_client.CoreV1Api(connection) + logs = await v1_api.read_namespaced_pod_log( + name=name, + namespace=namespace, + follow=False, + timestamps=True, + ) + logs = logs.splitlines() + for line in logs: + self.log.info("Container logs from %s", line) + return logs + except HTTPError: + self.log.exception("There was an error reading the kubernetes API.") + raise diff --git a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py index 954f3e018ebff..e52cfae870b89 100644 --- a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py @@ -317,14 +317,22 @@ def async_hook(self): ) @pytest.mark.asyncio - @mock.patch(GKE_STRING.format("Token"), mock.MagicMock()) + @pytest.mark.parametrize("mock_service_file", ("/tmp/service_file.json", None)) + @mock.patch(GKE_STRING.format("Token")) @mock.patch(GKE_STRING.format("GKEPodAsyncHook.get_conn")) @mock.patch(GKE_STRING.format("async_client.CoreV1Api.read_namespaced_pod")) - async def test_get_pod(self, read_namespace_pod_mock, get_conn_mock, async_hook): + async def test_get_pod( + self, read_namespace_pod_mock, get_conn_mock, mock_token, async_hook, mock_service_file + ): + async_hook.service_file_as_context = mock.MagicMock() + async_hook.service_file_as_context.return_value.__aenter__.return_value = mock_service_file + self.make_mock_awaitable(read_namespace_pod_mock) await async_hook.get_pod(name=POD_NAME, namespace=POD_NAMESPACE) - + mock_token.assert_called_with( + scopes=["https://www.googleapis.com/auth/cloud-platform"], service_file=mock_service_file + ) get_conn_mock.assert_called_once() read_namespace_pod_mock.assert_called_with( name=POD_NAME, @@ -332,14 +340,23 @@ async def test_get_pod(self, read_namespace_pod_mock, get_conn_mock, async_hook) ) @pytest.mark.asyncio - @mock.patch(GKE_STRING.format("Token"), mock.MagicMock()) + @pytest.mark.parametrize("mock_service_file", ("/tmp/service_file.json", None)) + @mock.patch(GKE_STRING.format("Token")) @mock.patch(GKE_STRING.format("GKEPodAsyncHook.get_conn")) @mock.patch(GKE_STRING.format("async_client.CoreV1Api.delete_namespaced_pod")) - async def test_delete_pod(self, delete_namespaced_pod, get_conn_mock, async_hook): + async def test_delete_pod( + self, delete_namespaced_pod, get_conn_mock, mock_token, async_hook, mock_service_file + ): + async_hook.service_file_as_context = mock.MagicMock() + async_hook.service_file_as_context.return_value.__aenter__.return_value = mock_service_file + self.make_mock_awaitable(delete_namespaced_pod) await async_hook.delete_pod(name=POD_NAME, namespace=POD_NAMESPACE) + mock_token.assert_called_with( + scopes=["https://www.googleapis.com/auth/cloud-platform"], service_file=mock_service_file + ) get_conn_mock.assert_called_once() delete_namespaced_pod.assert_called_with( name=POD_NAME, @@ -348,14 +365,23 @@ async def test_delete_pod(self, delete_namespaced_pod, get_conn_mock, async_hook ) @pytest.mark.asyncio - @mock.patch(GKE_STRING.format("Token"), mock.MagicMock()) + @pytest.mark.parametrize("mock_service_file", ("/tmp/service_file.json", None)) + @mock.patch(GKE_STRING.format("Token")) @mock.patch(GKE_STRING.format("GKEPodAsyncHook.get_conn")) @mock.patch(GKE_STRING.format("async_client.CoreV1Api.read_namespaced_pod_log")) - async def test_read_logs(self, read_namespaced_pod_log, get_conn_mock, async_hook, caplog): + async def test_read_logs( + self, read_namespaced_pod_log, get_conn_mock, mock_token, async_hook, mock_service_file, caplog + ): + async_hook.service_file_as_context = mock.MagicMock() + async_hook.service_file_as_context.return_value.__aenter__.return_value = mock_service_file + self.make_mock_awaitable(read_namespaced_pod_log, result="Test string #1\nTest string #2\n") await async_hook.read_logs(name=POD_NAME, namespace=POD_NAMESPACE) + mock_token.assert_called_with( + scopes=["https://www.googleapis.com/auth/cloud-platform"], service_file=mock_service_file + ) get_conn_mock.assert_called_once() read_namespaced_pod_log.assert_called_with( name=POD_NAME,