Skip to content

Commit

Permalink
Fix badly merged impersonation in GKEPodOperator (#19696)
Browse files Browse the repository at this point in the history
The #19518 was merged while we had false-positive test results
due to testing memory optmisation in CI - test failures went
unnoticed for the change.

This PR fixes the problem (both in tests and in the code) and
adds more tests to cover all scenarios
  • Loading branch information
potiuk committed Nov 19, 2021
1 parent 49b7e75 commit 3336bb6
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 2 deletions.
Expand Up @@ -365,7 +365,7 @@ def execute(self, context) -> Optional[str]:
if isinstance(self.impersonation_chain, str):
impersonation_account = self.impersonation_chain
elif len(self.impersonation_chain) == 1:
impersonation_account = self.impersonation_chain[:-1]
impersonation_account = self.impersonation_chain[0]
else:
raise AirflowException(
"Chained list of accounts is not supported, please specify only one service account"
Expand Down
74 changes: 73 additions & 1 deletion tests/providers/google/cloud/operators/test_kubernetes_engine.py
Expand Up @@ -313,7 +313,7 @@ def test_execute_with_impersonation_service_account(
type(file_mock.return_value.__enter__.return_value).name = PropertyMock(
side_effect=[FILE_NAME, '/path/to/new-file']
)
self.gke_op.impersonation_service_account = "test_account@example.com"
self.gke_op.impersonation_chain = "test_account@example.com"
self.gke_op.execute(None)

mock_gcp_hook.return_value.provide_authorized_gcloud.assert_called_once()
Expand All @@ -335,3 +335,75 @@ def test_execute_with_impersonation_service_account(
)

assert self.gke_op.config_file == FILE_NAME

@mock.patch.dict(os.environ, {})
@mock.patch(
"airflow.hooks.base.BaseHook.get_connections",
return_value=[
Connection(
extra=json.dumps(
{"extra__google_cloud_platform__keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}
)
)
],
)
@mock.patch('airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute')
@mock.patch('airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook')
@mock.patch('airflow.providers.google.cloud.operators.kubernetes_engine.execute_in_subprocess')
@mock.patch('tempfile.NamedTemporaryFile')
def test_execute_with_impersonation_service_chain_one_element(
self, file_mock, mock_execute_in_subprocess, mock_gcp_hook, exec_mock, get_con_mock
):
type(file_mock.return_value.__enter__.return_value).name = PropertyMock(
side_effect=[FILE_NAME, '/path/to/new-file']
)
self.gke_op.impersonation_chain = ["test_account@example.com"]
self.gke_op.execute(None)

mock_gcp_hook.return_value.provide_authorized_gcloud.assert_called_once()

mock_execute_in_subprocess.assert_called_once_with(
[
'gcloud',
'container',
'clusters',
'get-credentials',
CLUSTER_NAME,
'--zone',
PROJECT_LOCATION,
'--project',
TEST_GCP_PROJECT_ID,
'--impersonate-service-account',
'test_account@example.com',
]
)

assert self.gke_op.config_file == FILE_NAME

@mock.patch.dict(os.environ, {})
@mock.patch(
"airflow.hooks.base.BaseHook.get_connections",
return_value=[
Connection(
extra=json.dumps(
{"extra__google_cloud_platform__keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}
)
)
],
)
@mock.patch('airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute')
@mock.patch('airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook')
@mock.patch('airflow.providers.google.cloud.operators.kubernetes_engine.execute_in_subprocess')
@mock.patch('tempfile.NamedTemporaryFile')
def test_execute_with_impersonation_service_chain_more_elements(
self, file_mock, mock_execute_in_subprocess, mock_gcp_hook, exec_mock, get_con_mock
):
type(file_mock.return_value.__enter__.return_value).name = PropertyMock(
side_effect=[FILE_NAME, '/path/to/new-file']
)
self.gke_op.impersonation_chain = ["test_account@example.com", "test_account1@example.com"]
with pytest.raises(
AirflowException,
match="Chained list of accounts is not supported, please specify only one service account",
):
self.gke_op.execute(None)

0 comments on commit 3336bb6

Please sign in to comment.