class TestLifeSciencesHookWithoutProjectId(unittest.TestCase): def setUp(self): with mock.patch( "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__", new=mock_base_gcp_hook_no_default_project_id, ): self.hook = LifeSciencesHook(gcp_conn_id="test") @mock.patch("airflow.providers.google.cloud.hooks.life_sciences.LifeSciencesHook._authorize") @mock.patch("airflow.providers.google.cloud.hooks.life_sciences.build") def test_life_science_client_creation(self, mock_build, mock_authorize): result = self.hook.get_conn() mock_build.assert_called_once_with( 'lifesciences', 'v2beta', http=mock_authorize.return_value, cache_discovery=False ) self.assertEqual(mock_build.return_value, result) self.assertEqual(self.hook._conn, result) @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, return_value=None, ) @mock.patch("airflow.providers.google.cloud.hooks.life_sciences.LifeSciencesHook.get_conn") def test_run_pipeline(self, get_conn_mock, mock_project_id): # pylint: disable=unused-argument with self.assertRaises(AirflowException) as e: self.hook.run_pipeline(body={}, location=TEST_LOCATION) # pylint: disable=no-value-for-parameter self.assertEqual( "The project id must be passed either as keyword project_id parameter or as project_id extra in " "Google Cloud connection definition. Both are not set!", str(e.exception), )
def execute(self, context): hook = LifeSciencesHook(gcp_conn_id=self.gcp_conn_id, api_version=self.api_version) return hook.run_pipeline(body=self.body, location=self.location, project_id=self.project_id)
def execute(self, context) -> dict: hook = LifeSciencesHook( gcp_conn_id=self.gcp_conn_id, api_version=self.api_version, impersonation_chain=self.impersonation_chain, ) return hook.run_pipeline(body=self.body, location=self.location, project_id=self.project_id)
class TestLifeSciencesHookWithPassedProjectId(unittest.TestCase): def setUp(self): with mock.patch( "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__", new=mock_base_gcp_hook_default_project_id, ): self.hook = LifeSciencesHook(gcp_conn_id="test") def test_location_path(self): path = 'projects/life-science-project-id/locations/test-location' path2 = self.hook._location_path(project_id=TEST_PROJECT_ID, location=TEST_LOCATION) self.assertEqual(path, path2) @mock.patch( "airflow.providers.google.cloud.hooks.life_sciences.LifeSciencesHook._authorize" ) @mock.patch("airflow.providers.google.cloud.hooks.life_sciences.build") def test_life_science_client_creation(self, mock_build, mock_authorize): result = self.hook.get_conn() mock_build.assert_called_once_with('lifesciences', 'v2beta', http=mock_authorize.return_value, cache_discovery=False) self.assertEqual(mock_build.return_value, result) self.assertEqual(self.hook._conn, result) @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) @mock.patch( "airflow.providers.google.cloud.hooks.life_sciences.LifeSciencesHook.get_conn" ) def test_run_pipeline_immediately_complete(self, get_conn_mock, mock_project_id): service_mock = get_conn_mock.return_value # fmt: off service_mock.projects.return_value \ .locations.return_value \ .pipelines.return_value \ .run.return_value \ .execute.return_value = TEST_OPERATION service_mock.projects.return_value \ .locations.return_value \ .operations.return_value \ .get.return_value \ .execute.return_value = TEST_DONE_OPERATION result = self.hook.run_pipeline( body={}, # pylint: disable=no-value-for-parameter location=TEST_LOCATION) parent = self.hook. \ _location_path(location=TEST_LOCATION) # pylint: disable=no-value-for-parameter service_mock.projects.return_value.locations.return_value \ .pipelines.return_value.run \ .assert_called_once_with(body={}, parent=parent) # fmt: on self.assertEqual(result, TEST_OPERATION) @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) @mock.patch( "airflow.providers.google.cloud.hooks.life_sciences.LifeSciencesHook.get_conn" ) @mock.patch("airflow.providers.google.cloud.hooks.life_sciences.time.sleep" ) def test_waiting_operation(self, _, get_conn_mock, mock_project_id): service_mock = get_conn_mock.return_value # fmt: off service_mock.projects.return_value \ .locations.return_value \ .pipelines.return_value \ .run.return_value \ .execute.return_value = TEST_OPERATION execute_mock = mock.Mock( **{"side_effect": [TEST_WAITING_OPERATION, TEST_DONE_OPERATION]}) service_mock.projects.return_value \ .locations.return_value \ .operations.return_value \ .get.return_value \ .execute = execute_mock # fmt: on result = self.hook.run_pipeline(body={}, location=TEST_LOCATION, project_id=TEST_PROJECT_ID) self.assertEqual(result, TEST_OPERATION) @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) @mock.patch( "airflow.providers.google.cloud.hooks.life_sciences.LifeSciencesHook.get_conn" ) @mock.patch("airflow.providers.google.cloud.hooks.life_sciences.time.sleep" ) def test_error_operation(self, _, get_conn_mock, mock_project_id): service_mock = get_conn_mock.return_value # fmt: off service_mock.projects.return_value \ .locations.return_value \ .pipelines.return_value \ .run.return_value \ .execute.return_value = TEST_OPERATION execute_mock = mock.Mock( **{"side_effect": [TEST_WAITING_OPERATION, TEST_ERROR_OPERATION]}) service_mock.projects.return_value \ .locations.return_value \ .operations.return_value \ .get.return_value \ .execute = execute_mock # fmt: on with self.assertRaisesRegex(AirflowException, "error"): self.hook.run_pipeline(body={}, location=TEST_LOCATION, project_id=TEST_PROJECT_ID)