From e1f96415e5632a9f97ae0e100287487cb8e9f1a2 Mon Sep 17 00:00:00 2001 From: Jonas Grabber Date: Mon, 20 Dec 2021 17:27:02 +0100 Subject: [PATCH 1/4] Fix setting of project ID in provide_authorized_gcloud --- airflow/providers/google/common/hooks/base_google.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py index a149831cab5e2..ca7fb2c717c38 100644 --- a/airflow/providers/google/common/hooks/base_google.py +++ b/airflow/providers/google/common/hooks/base_google.py @@ -515,9 +515,6 @@ def provide_authorized_gcloud(self): f"--key-file={os.environ[CREDENTIALS]}", ] ) - if project_id: - # Don't display stdout/stderr for security reason - check_output(["gcloud", "config", "set", "core/project", project_id]) elif os.path.exists(credentials_path): # If we are logged in by `gcloud auth application-default` then we need to log in manually. # This will make the `gcloud auth application-default` and `gcloud auth` credentials equals. @@ -539,6 +536,11 @@ def provide_authorized_gcloud(self): creds_content["refresh_token"], ] ) + + if project_id: + # Don't display stdout/stderr for security reason + check_output(["gcloud", "config", "set", "core/project", project_id]) + yield @staticmethod From 6941e5fa28b2708b41acd03bfbc22c215ecb1a56 Mon Sep 17 00:00:00 2001 From: Jonas Grabber Date: Tue, 21 Dec 2021 12:39:51 +0100 Subject: [PATCH 2/4] Fix tests in test_base_google Some tests use has_calls instead of assert_has_calls, but don't assert on the returned value.. --- airflow/providers/google/common/hooks/base_google.py | 2 +- tests/providers/airbyte/hooks/test_airbyte.py | 10 +++++----- .../google/cloud/hooks/test_cloud_memorystore.py | 6 +++--- tests/providers/google/cloud/hooks/test_dataproc.py | 4 ++-- tests/providers/google/cloud/hooks/test_pubsub.py | 2 +- .../google/cloud/transfers/test_sheets_to_gcs.py | 4 ++-- .../providers/google/common/hooks/test_base_google.py | 4 ++-- tests/providers/google/suite/operators/test_sheets.py | 2 +- tests/providers/http/operators/test_http.py | 2 +- 9 files changed, 18 insertions(+), 18 deletions(-) diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py index ca7fb2c717c38..0faaddf899549 100644 --- a/airflow/providers/google/common/hooks/base_google.py +++ b/airflow/providers/google/common/hooks/base_google.py @@ -540,7 +540,7 @@ def provide_authorized_gcloud(self): if project_id: # Don't display stdout/stderr for security reason check_output(["gcloud", "config", "set", "core/project", project_id]) - + yield @staticmethod diff --git a/tests/providers/airbyte/hooks/test_airbyte.py b/tests/providers/airbyte/hooks/test_airbyte.py index f72d53847ce5a..eff71f9d77bcd 100644 --- a/tests/providers/airbyte/hooks/test_airbyte.py +++ b/tests/providers/airbyte/hooks/test_airbyte.py @@ -87,7 +87,7 @@ def test_wait_for_job_error(self, mock_get_job): self.hook.wait_for_job(job_id=self.job_id, wait_seconds=0) calls = [mock.call(job_id=self.job_id), mock.call(job_id=self.job_id)] - assert mock_get_job.has_calls(calls) + assert mock_get_job.assert_has_calls(calls) @mock.patch('airflow.providers.airbyte.hooks.airbyte.AirbyteHook.get_job') def test_wait_for_job_incomplete_succeeded(self, mock_get_job): @@ -98,7 +98,7 @@ def test_wait_for_job_incomplete_succeeded(self, mock_get_job): self.hook.wait_for_job(job_id=self.job_id, wait_seconds=0) calls = [mock.call(job_id=self.job_id), mock.call(job_id=self.job_id)] - assert mock_get_job.has_calls(calls) + assert mock_get_job.assert_has_calls(calls) @mock.patch('airflow.providers.airbyte.hooks.airbyte.AirbyteHook.get_job') def test_wait_for_job_timeout(self, mock_get_job): @@ -111,7 +111,7 @@ def test_wait_for_job_timeout(self, mock_get_job): self.hook.wait_for_job(job_id=self.job_id, wait_seconds=2, timeout=1) calls = [mock.call(job_id=self.job_id), mock.call(job_id=self.job_id), mock.call(job_id=self.job_id)] - assert mock_get_job.has_calls(calls) + assert mock_get_job.assert_has_calls(calls) @mock.patch('airflow.providers.airbyte.hooks.airbyte.AirbyteHook.get_job') def test_wait_for_job_state_unrecognized(self, mock_get_job): @@ -123,7 +123,7 @@ def test_wait_for_job_state_unrecognized(self, mock_get_job): self.hook.wait_for_job(job_id=self.job_id, wait_seconds=0) calls = [mock.call(job_id=self.job_id), mock.call(job_id=self.job_id)] - assert mock_get_job.has_calls(calls) + assert mock_get_job.assert_has_calls(calls) @mock.patch('airflow.providers.airbyte.hooks.airbyte.AirbyteHook.get_job') def test_wait_for_job_cancelled(self, mock_get_job): @@ -135,7 +135,7 @@ def test_wait_for_job_cancelled(self, mock_get_job): self.hook.wait_for_job(job_id=self.job_id, wait_seconds=0) calls = [mock.call(job_id=self.job_id), mock.call(job_id=self.job_id)] - assert mock_get_job.has_calls(calls) + assert mock_get_job.assert_has_calls(calls) @requests_mock.mock() def test_connection_success(self, m): diff --git a/tests/providers/google/cloud/hooks/test_cloud_memorystore.py b/tests/providers/google/cloud/hooks/test_cloud_memorystore.py index 98f3b10ff37f9..bbee3448852cb 100644 --- a/tests/providers/google/cloud/hooks/test_cloud_memorystore.py +++ b/tests/providers/google/cloud/hooks/test_cloud_memorystore.py @@ -112,7 +112,7 @@ def test_create_instance_when_not_exists(self, mock_get_conn, mock_project_id): timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) - mock_get_conn.return_value.get_instance.has_calls( + mock_get_conn.return_value.get_instance.assert_has_calls( [ mock.call(name=TEST_NAME, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA), mock.call(name=TEST_NAME, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA), @@ -266,7 +266,7 @@ def test_create_instance_when_not_exists(self, mock_get_conn): timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) - mock_get_conn.return_value.get_instance.has_calls( + mock_get_conn.return_value.get_instance.assert_has_calls( [ mock.call( name="projects/test-project-id/locations/test-location/instances/test-instance-id", @@ -503,7 +503,7 @@ def test_create_instance_when_not_exists(self, mock_get_conn, mock_project_id): timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) - mock_get_conn.return_value.get_instance.has_calls( + mock_get_conn.return_value.get_instance.assert_has_calls( [ mock.call(name=TEST_NAME, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA), mock.call(name=TEST_NAME, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA), diff --git a/tests/providers/google/cloud/hooks/test_dataproc.py b/tests/providers/google/cloud/hooks/test_dataproc.py index 76714b4e00678..733b725ec5437 100644 --- a/tests/providers/google/cloud/hooks/test_dataproc.py +++ b/tests/providers/google/cloud/hooks/test_dataproc.py @@ -493,7 +493,7 @@ def test_wait_for_job(self, mock_get_job): mock.call(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT), mock.call(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT), ] - mock_get_job.has_calls(calls) + mock_get_job.assert_has_calls(calls) @mock.patch(DATAPROC_STRING.format("DataprocHook.get_job")) def test_wait_for_job_deprecation_warning(self, mock_get_job): @@ -514,7 +514,7 @@ def test_wait_for_job_deprecation_warning(self, mock_get_job): mock.call(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT), mock.call(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT), ] - mock_get_job.has_calls(calls) + mock_get_job.assert_has_calls(calls) assert warning_message == str(warnings[0].message) with pytest.raises(TypeError): diff --git a/tests/providers/google/cloud/hooks/test_pubsub.py b/tests/providers/google/cloud/hooks/test_pubsub.py index 027c2b2f3d89e..36a1c17f27fb6 100644 --- a/tests/providers/google/cloud/hooks/test_pubsub.py +++ b/tests/providers/google/cloud/hooks/test_pubsub.py @@ -380,7 +380,7 @@ def test_publish(self, mock_service): mock.call(topic=EXPANDED_TOPIC, data=message.get("data", b''), **message.get('attributes', {})) for message in TEST_MESSAGES ] - publish_method.has_calls(calls) + publish_method.assert_has_calls(calls) @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn')) def test_publish_api_call_error(self, mock_service): diff --git a/tests/providers/google/cloud/transfers/test_sheets_to_gcs.py b/tests/providers/google/cloud/transfers/test_sheets_to_gcs.py index fba432d95f9c7..1e31a0adc810d 100644 --- a/tests/providers/google/cloud/transfers/test_sheets_to_gcs.py +++ b/tests/providers/google/cloud/transfers/test_sheets_to_gcs.py @@ -114,9 +114,9 @@ def test_execute(self, mock_upload_data, mock_xcom, mock_sheet_hook, mock_gcs_ho ) calls = [mock.call(spreadsheet_id=SPREADSHEET_ID, range_=r) for r in RANGES] - mock_sheet_hook.return_value.get_values.has_calls(calls) + mock_sheet_hook.return_value.get_values.assert_has_calls(calls) calls = [mock.call(mock_gcs_hook, mock_sheet_hook, r, v) for r, v in zip(RANGES, data)] - mock_upload_data.has_calls(calls) + mock_upload_data.assert_has_calls(calls) mock_xcom.assert_called_once_with(context, "destination_objects", [PATH, PATH]) diff --git a/tests/providers/google/common/hooks/test_base_google.py b/tests/providers/google/common/hooks/test_base_google.py index fd4cd613bb3cb..93b3b2d2bfd0c 100644 --- a/tests/providers/google/common/hooks/test_base_google.py +++ b/tests/providers/google/common/hooks/test_base_google.py @@ -742,12 +742,12 @@ def test_provide_authorized_gcloud_via_gcloud_application_default( # Do nothing pass - mock_check_output.has_calls( + mock_check_output.assert_has_calls( [ mock.call(['gcloud', 'config', 'set', 'auth/client_id', 'CLIENT_ID']), mock.call(['gcloud', 'config', 'set', 'auth/client_secret', 'CLIENT_SECRET']), - mock.call(['gcloud', 'config', 'set', 'core/project', 'PROJECT_ID']), mock.call(['gcloud', 'auth', 'activate-refresh-token', 'CLIENT_ID', 'REFRESH_TOKEN']), + mock.call(['gcloud', 'config', 'set', 'core/project', 'PROJECT_ID']), ], any_order=False, ) diff --git a/tests/providers/google/suite/operators/test_sheets.py b/tests/providers/google/suite/operators/test_sheets.py index 0bf2da4fffe07..b1c741cd5a942 100644 --- a/tests/providers/google/suite/operators/test_sheets.py +++ b/tests/providers/google/suite/operators/test_sheets.py @@ -47,4 +47,4 @@ def test_execute(self, mock_xcom, mock_hook): mock.call(context, "spreadsheet_id", SPREADSHEET_ID), mock.call(context, "spreadsheet_url", SPREADSHEET_URL), ] - mock_xcom.has_calls(calls) + mock_xcom.assert_has_calls(calls) diff --git a/tests/providers/http/operators/test_http.py b/tests/providers/http/operators/test_http.py index d4c622c29e6a5..a4caef18b2127 100644 --- a/tests/providers/http/operators/test_http.py +++ b/tests/providers/http/operators/test_http.py @@ -47,7 +47,7 @@ def test_response_in_logs(self, m): with mock.patch.object(operator.log, 'info') as mock_info: operator.execute(None) calls = [mock.call('Example.com fake response'), mock.call('Example.com fake response')] - mock_info.has_calls(calls) + mock_info.assert_has_calls(calls) @requests_mock.mock() def test_response_in_logs_after_failed_check(self, m): From d29f2c37d47fe278ef4d26dfb5a653b94cf3d81c Mon Sep 17 00:00:00 2001 From: Jonas Grabber Date: Tue, 21 Dec 2021 13:49:27 +0100 Subject: [PATCH 3/4] Fix formatting --- airflow/providers/google/common/hooks/base_google.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py index 0faaddf899549..0e2c6b17ad567 100644 --- a/airflow/providers/google/common/hooks/base_google.py +++ b/airflow/providers/google/common/hooks/base_google.py @@ -536,7 +536,7 @@ def provide_authorized_gcloud(self): creds_content["refresh_token"], ] ) - + if project_id: # Don't display stdout/stderr for security reason check_output(["gcloud", "config", "set", "core/project", project_id]) From 82608a4fdeabd29ccb0850f58f601f95c82030f4 Mon Sep 17 00:00:00 2001 From: Jonas Grabber Date: Tue, 21 Dec 2021 14:50:05 +0100 Subject: [PATCH 4/4] Revert fixing tests other than for base_google.py This partially reverts commit 6941e5fa28b2708b41acd03bfbc22c215ecb1a56. --- tests/providers/airbyte/hooks/test_airbyte.py | 10 +++++----- .../google/cloud/hooks/test_cloud_memorystore.py | 6 +++--- tests/providers/google/cloud/hooks/test_dataproc.py | 4 ++-- tests/providers/google/cloud/hooks/test_pubsub.py | 2 +- .../google/cloud/transfers/test_sheets_to_gcs.py | 4 ++-- tests/providers/google/suite/operators/test_sheets.py | 2 +- tests/providers/http/operators/test_http.py | 2 +- 7 files changed, 15 insertions(+), 15 deletions(-) diff --git a/tests/providers/airbyte/hooks/test_airbyte.py b/tests/providers/airbyte/hooks/test_airbyte.py index eff71f9d77bcd..f72d53847ce5a 100644 --- a/tests/providers/airbyte/hooks/test_airbyte.py +++ b/tests/providers/airbyte/hooks/test_airbyte.py @@ -87,7 +87,7 @@ def test_wait_for_job_error(self, mock_get_job): self.hook.wait_for_job(job_id=self.job_id, wait_seconds=0) calls = [mock.call(job_id=self.job_id), mock.call(job_id=self.job_id)] - assert mock_get_job.assert_has_calls(calls) + assert mock_get_job.has_calls(calls) @mock.patch('airflow.providers.airbyte.hooks.airbyte.AirbyteHook.get_job') def test_wait_for_job_incomplete_succeeded(self, mock_get_job): @@ -98,7 +98,7 @@ def test_wait_for_job_incomplete_succeeded(self, mock_get_job): self.hook.wait_for_job(job_id=self.job_id, wait_seconds=0) calls = [mock.call(job_id=self.job_id), mock.call(job_id=self.job_id)] - assert mock_get_job.assert_has_calls(calls) + assert mock_get_job.has_calls(calls) @mock.patch('airflow.providers.airbyte.hooks.airbyte.AirbyteHook.get_job') def test_wait_for_job_timeout(self, mock_get_job): @@ -111,7 +111,7 @@ def test_wait_for_job_timeout(self, mock_get_job): self.hook.wait_for_job(job_id=self.job_id, wait_seconds=2, timeout=1) calls = [mock.call(job_id=self.job_id), mock.call(job_id=self.job_id), mock.call(job_id=self.job_id)] - assert mock_get_job.assert_has_calls(calls) + assert mock_get_job.has_calls(calls) @mock.patch('airflow.providers.airbyte.hooks.airbyte.AirbyteHook.get_job') def test_wait_for_job_state_unrecognized(self, mock_get_job): @@ -123,7 +123,7 @@ def test_wait_for_job_state_unrecognized(self, mock_get_job): self.hook.wait_for_job(job_id=self.job_id, wait_seconds=0) calls = [mock.call(job_id=self.job_id), mock.call(job_id=self.job_id)] - assert mock_get_job.assert_has_calls(calls) + assert mock_get_job.has_calls(calls) @mock.patch('airflow.providers.airbyte.hooks.airbyte.AirbyteHook.get_job') def test_wait_for_job_cancelled(self, mock_get_job): @@ -135,7 +135,7 @@ def test_wait_for_job_cancelled(self, mock_get_job): self.hook.wait_for_job(job_id=self.job_id, wait_seconds=0) calls = [mock.call(job_id=self.job_id), mock.call(job_id=self.job_id)] - assert mock_get_job.assert_has_calls(calls) + assert mock_get_job.has_calls(calls) @requests_mock.mock() def test_connection_success(self, m): diff --git a/tests/providers/google/cloud/hooks/test_cloud_memorystore.py b/tests/providers/google/cloud/hooks/test_cloud_memorystore.py index bbee3448852cb..98f3b10ff37f9 100644 --- a/tests/providers/google/cloud/hooks/test_cloud_memorystore.py +++ b/tests/providers/google/cloud/hooks/test_cloud_memorystore.py @@ -112,7 +112,7 @@ def test_create_instance_when_not_exists(self, mock_get_conn, mock_project_id): timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) - mock_get_conn.return_value.get_instance.assert_has_calls( + mock_get_conn.return_value.get_instance.has_calls( [ mock.call(name=TEST_NAME, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA), mock.call(name=TEST_NAME, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA), @@ -266,7 +266,7 @@ def test_create_instance_when_not_exists(self, mock_get_conn): timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) - mock_get_conn.return_value.get_instance.assert_has_calls( + mock_get_conn.return_value.get_instance.has_calls( [ mock.call( name="projects/test-project-id/locations/test-location/instances/test-instance-id", @@ -503,7 +503,7 @@ def test_create_instance_when_not_exists(self, mock_get_conn, mock_project_id): timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) - mock_get_conn.return_value.get_instance.assert_has_calls( + mock_get_conn.return_value.get_instance.has_calls( [ mock.call(name=TEST_NAME, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA), mock.call(name=TEST_NAME, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA), diff --git a/tests/providers/google/cloud/hooks/test_dataproc.py b/tests/providers/google/cloud/hooks/test_dataproc.py index 733b725ec5437..76714b4e00678 100644 --- a/tests/providers/google/cloud/hooks/test_dataproc.py +++ b/tests/providers/google/cloud/hooks/test_dataproc.py @@ -493,7 +493,7 @@ def test_wait_for_job(self, mock_get_job): mock.call(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT), mock.call(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT), ] - mock_get_job.assert_has_calls(calls) + mock_get_job.has_calls(calls) @mock.patch(DATAPROC_STRING.format("DataprocHook.get_job")) def test_wait_for_job_deprecation_warning(self, mock_get_job): @@ -514,7 +514,7 @@ def test_wait_for_job_deprecation_warning(self, mock_get_job): mock.call(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT), mock.call(region=GCP_LOCATION, job_id=JOB_ID, project_id=GCP_PROJECT), ] - mock_get_job.assert_has_calls(calls) + mock_get_job.has_calls(calls) assert warning_message == str(warnings[0].message) with pytest.raises(TypeError): diff --git a/tests/providers/google/cloud/hooks/test_pubsub.py b/tests/providers/google/cloud/hooks/test_pubsub.py index 36a1c17f27fb6..027c2b2f3d89e 100644 --- a/tests/providers/google/cloud/hooks/test_pubsub.py +++ b/tests/providers/google/cloud/hooks/test_pubsub.py @@ -380,7 +380,7 @@ def test_publish(self, mock_service): mock.call(topic=EXPANDED_TOPIC, data=message.get("data", b''), **message.get('attributes', {})) for message in TEST_MESSAGES ] - publish_method.assert_has_calls(calls) + publish_method.has_calls(calls) @mock.patch(PUBSUB_STRING.format('PubSubHook.get_conn')) def test_publish_api_call_error(self, mock_service): diff --git a/tests/providers/google/cloud/transfers/test_sheets_to_gcs.py b/tests/providers/google/cloud/transfers/test_sheets_to_gcs.py index 1e31a0adc810d..fba432d95f9c7 100644 --- a/tests/providers/google/cloud/transfers/test_sheets_to_gcs.py +++ b/tests/providers/google/cloud/transfers/test_sheets_to_gcs.py @@ -114,9 +114,9 @@ def test_execute(self, mock_upload_data, mock_xcom, mock_sheet_hook, mock_gcs_ho ) calls = [mock.call(spreadsheet_id=SPREADSHEET_ID, range_=r) for r in RANGES] - mock_sheet_hook.return_value.get_values.assert_has_calls(calls) + mock_sheet_hook.return_value.get_values.has_calls(calls) calls = [mock.call(mock_gcs_hook, mock_sheet_hook, r, v) for r, v in zip(RANGES, data)] - mock_upload_data.assert_has_calls(calls) + mock_upload_data.has_calls(calls) mock_xcom.assert_called_once_with(context, "destination_objects", [PATH, PATH]) diff --git a/tests/providers/google/suite/operators/test_sheets.py b/tests/providers/google/suite/operators/test_sheets.py index b1c741cd5a942..0bf2da4fffe07 100644 --- a/tests/providers/google/suite/operators/test_sheets.py +++ b/tests/providers/google/suite/operators/test_sheets.py @@ -47,4 +47,4 @@ def test_execute(self, mock_xcom, mock_hook): mock.call(context, "spreadsheet_id", SPREADSHEET_ID), mock.call(context, "spreadsheet_url", SPREADSHEET_URL), ] - mock_xcom.assert_has_calls(calls) + mock_xcom.has_calls(calls) diff --git a/tests/providers/http/operators/test_http.py b/tests/providers/http/operators/test_http.py index a4caef18b2127..d4c622c29e6a5 100644 --- a/tests/providers/http/operators/test_http.py +++ b/tests/providers/http/operators/test_http.py @@ -47,7 +47,7 @@ def test_response_in_logs(self, m): with mock.patch.object(operator.log, 'info') as mock_info: operator.execute(None) calls = [mock.call('Example.com fake response'), mock.call('Example.com fake response')] - mock_info.assert_has_calls(calls) + mock_info.has_calls(calls) @requests_mock.mock() def test_response_in_logs_after_failed_check(self, m):