From f8c094ad4b6b23cd23a614eb7c9a777a4a102de0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Freyr=20S=C3=A6varsson?= Date: Tue, 6 Feb 2024 13:42:38 +0000 Subject: [PATCH 1/5] Add project_id to hook.get_job calls --- airflow/providers/google/cloud/operators/cloud_run.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/providers/google/cloud/operators/cloud_run.py b/airflow/providers/google/cloud/operators/cloud_run.py index 7aab64096f059..695ea6708f92c 100644 --- a/airflow/providers/google/cloud/operators/cloud_run.py +++ b/airflow/providers/google/cloud/operators/cloud_run.py @@ -302,7 +302,7 @@ def execute(self, context: Context): if not self.deferrable: result: Execution = self._wait_for_operation(self.operation) self._fail_if_execution_failed(result) - job = hook.get_job(job_name=result.job, region=self.region) + job = hook.get_job(job_name=result.job, region=self.region, project_id=self.project_id) return Job.to_dict(job) else: self.defer( @@ -333,7 +333,7 @@ def execute_complete(self, context: Context, event: dict): hook: CloudRunHook = CloudRunHook(self.gcp_conn_id, self.impersonation_chain) - job = hook.get_job(job_name=event["job_name"], region=self.region) + job = hook.get_job(job_name=event["job_name"], region=self.region, project_id=self.project_id) return Job.to_dict(job) def _fail_if_execution_failed(self, execution: Execution): From 5bdf8ed948fd56c9f9defdc84277b29ce2dbc54c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Freyr=20S=C3=A6varsson?= Date: Wed, 7 Feb 2024 10:47:25 +0000 Subject: [PATCH 2/5] Update tests to verify that get_job is called with all necessary parameters. --- .../google/cloud/operators/test_cloud_run.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/providers/google/cloud/operators/test_cloud_run.py b/tests/providers/google/cloud/operators/test_cloud_run.py index 459dd81b27067..ecbee4ff4eb7c 100644 --- a/tests/providers/google/cloud/operators/test_cloud_run.py +++ b/tests/providers/google/cloud/operators/test_cloud_run.py @@ -102,6 +102,10 @@ def test_execute_success(self, hook_mock): operator.execute(context=mock.MagicMock()) + hook_mock.return_value.get_job.assert_called_once_with( + job_name=JOB_NAME, region=REGION, project_id=PROJECT_ID + ) + hook_mock.return_value.execute_job.assert_called_once_with( job_name=JOB_NAME, region=REGION, project_id=PROJECT_ID, overrides=None ) @@ -214,6 +218,10 @@ def test_execute_deferrable_execute_complete_method_success(self, hook_mock): event = {"status": RunJobStatus.SUCCESS.value, "job_name": JOB_NAME} result = operator.execute_complete(mock.MagicMock(), event) + + hook_mock.return_value.get_job.assert_called_once_with( + job_name=JOB_NAME, region=REGION, project_id=PROJECT_ID + ) assert result["name"] == JOB_NAME @mock.patch(CLOUD_RUN_HOOK_PATH) @@ -233,6 +241,10 @@ def test_execute_overrides(self, hook_mock): operator.execute(context=mock.MagicMock()) + hook_mock.return_value.get_job.assert_called_once_with( + job_name=JOB_NAME, region=REGION, project_id=PROJECT_ID + ) + hook_mock.return_value.execute_job.assert_called_once_with( job_name=JOB_NAME, region=REGION, project_id=PROJECT_ID, overrides=overrides ) From bcb8c68403e383faff7725aa94bb82c8b48389c2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Freyr=20S=C3=A6varsson?= Date: Fri, 9 Feb 2024 09:33:25 +0000 Subject: [PATCH 3/5] remove whitespace --- tests/providers/google/cloud/operators/test_cloud_run.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/providers/google/cloud/operators/test_cloud_run.py b/tests/providers/google/cloud/operators/test_cloud_run.py index ecbee4ff4eb7c..e2ddb490a67aa 100644 --- a/tests/providers/google/cloud/operators/test_cloud_run.py +++ b/tests/providers/google/cloud/operators/test_cloud_run.py @@ -105,7 +105,7 @@ def test_execute_success(self, hook_mock): hook_mock.return_value.get_job.assert_called_once_with( job_name=JOB_NAME, region=REGION, project_id=PROJECT_ID ) - + hook_mock.return_value.execute_job.assert_called_once_with( job_name=JOB_NAME, region=REGION, project_id=PROJECT_ID, overrides=None ) @@ -218,7 +218,7 @@ def test_execute_deferrable_execute_complete_method_success(self, hook_mock): event = {"status": RunJobStatus.SUCCESS.value, "job_name": JOB_NAME} result = operator.execute_complete(mock.MagicMock(), event) - + hook_mock.return_value.get_job.assert_called_once_with( job_name=JOB_NAME, region=REGION, project_id=PROJECT_ID ) @@ -244,7 +244,7 @@ def test_execute_overrides(self, hook_mock): hook_mock.return_value.get_job.assert_called_once_with( job_name=JOB_NAME, region=REGION, project_id=PROJECT_ID ) - + hook_mock.return_value.execute_job.assert_called_once_with( job_name=JOB_NAME, region=REGION, project_id=PROJECT_ID, overrides=overrides ) From ce1890d259c25d5917528e1b2e57f6645c6c49d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Freyr=20S=C3=A6varsson?= Date: Mon, 12 Feb 2024 11:22:01 +0000 Subject: [PATCH 4/5] Update test_cloud_run.py --- tests/providers/google/cloud/operators/test_cloud_run.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/providers/google/cloud/operators/test_cloud_run.py b/tests/providers/google/cloud/operators/test_cloud_run.py index e2ddb490a67aa..0e160fed7b09b 100644 --- a/tests/providers/google/cloud/operators/test_cloud_run.py +++ b/tests/providers/google/cloud/operators/test_cloud_run.py @@ -103,7 +103,7 @@ def test_execute_success(self, hook_mock): operator.execute(context=mock.MagicMock()) hook_mock.return_value.get_job.assert_called_once_with( - job_name=JOB_NAME, region=REGION, project_id=PROJECT_ID + region=REGION, project_id=PROJECT_ID ) hook_mock.return_value.execute_job.assert_called_once_with( @@ -220,7 +220,7 @@ def test_execute_deferrable_execute_complete_method_success(self, hook_mock): result = operator.execute_complete(mock.MagicMock(), event) hook_mock.return_value.get_job.assert_called_once_with( - job_name=JOB_NAME, region=REGION, project_id=PROJECT_ID + region=REGION, project_id=PROJECT_ID ) assert result["name"] == JOB_NAME @@ -242,7 +242,7 @@ def test_execute_overrides(self, hook_mock): operator.execute(context=mock.MagicMock()) hook_mock.return_value.get_job.assert_called_once_with( - job_name=JOB_NAME, region=REGION, project_id=PROJECT_ID + region=REGION, project_id=PROJECT_ID ) hook_mock.return_value.execute_job.assert_called_once_with( From f81f88ee990791e272bb024853ae03d0290c6456 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Freyr=20S=C3=A6varsson?= Date: Mon, 12 Feb 2024 11:32:46 +0000 Subject: [PATCH 5/5] Update test_cloud_run.py --- tests/providers/google/cloud/operators/test_cloud_run.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/providers/google/cloud/operators/test_cloud_run.py b/tests/providers/google/cloud/operators/test_cloud_run.py index 0e160fed7b09b..855190efd149c 100644 --- a/tests/providers/google/cloud/operators/test_cloud_run.py +++ b/tests/providers/google/cloud/operators/test_cloud_run.py @@ -103,7 +103,7 @@ def test_execute_success(self, hook_mock): operator.execute(context=mock.MagicMock()) hook_mock.return_value.get_job.assert_called_once_with( - region=REGION, project_id=PROJECT_ID + job_name=mock.ANY, region=REGION, project_id=PROJECT_ID ) hook_mock.return_value.execute_job.assert_called_once_with( @@ -220,7 +220,7 @@ def test_execute_deferrable_execute_complete_method_success(self, hook_mock): result = operator.execute_complete(mock.MagicMock(), event) hook_mock.return_value.get_job.assert_called_once_with( - region=REGION, project_id=PROJECT_ID + job_name=mock.ANY, region=REGION, project_id=PROJECT_ID ) assert result["name"] == JOB_NAME @@ -242,7 +242,7 @@ def test_execute_overrides(self, hook_mock): operator.execute(context=mock.MagicMock()) hook_mock.return_value.get_job.assert_called_once_with( - region=REGION, project_id=PROJECT_ID + job_name=mock.ANY, region=REGION, project_id=PROJECT_ID ) hook_mock.return_value.execute_job.assert_called_once_with(