def execute(self, context: Dict): hook = CloudMemorystoreHook(gcp_conn_id=self.gcp_conn_id) result = hook.create_instance( location=self.location, instance_id=self.instance_id, instance=self.instance, project_id=self.project_id, retry=self.retry, timeout=self.timeout, metadata=self.metadata, ) return MessageToDict(result)
def execute(self, context: Dict): hook = CloudMemorystoreHook(gcp_conn_id=self.gcp_conn_id) hook.create_instance( location=self.location, instance_id=self.instance_id, instance=self.instance, project_id=self.project_id, retry=self.retry, timeout=self.timeout, metadata=self.metadata, ) hook.import_instance( location=self.location, instance=self.instance, input_config=self.input_config, project_id=self.project_id, retry=self.retry, timeout=self.timeout, metadata=self.metadata, )
class TestCloudMemorystoreWithDefaultProjectIdHook(TestCase): def setUp(self,): with mock.patch( "airflow.gcp.hooks.cloud_memorystore.CloudMemorystoreHook.__init__", new=mock_base_gcp_hook_default_project_id, ): self.hook = CloudMemorystoreHook(gcp_conn_id="test") @mock.patch( 'airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST ) @mock.patch( # type: ignore "airflow.gcp.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn" ) def test_create_instance_when_exists(self, mock_get_conn, mock_project_id): mock_get_conn.return_value.get_instance.return_value = Instance(name=TEST_NAME) result = self.hook.create_instance( location=TEST_LOCATION, instance_id=TEST_INSTANCE_ID, instance=Instance(name=TEST_NAME), retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) mock_get_conn.return_value.get_instance.assert_called_once_with( name=TEST_NAME_DEFAULT_PROJECT_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA ) self.assertEqual(Instance(name=TEST_NAME), result) @mock.patch( 'airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST ) @mock.patch( # type: ignore "airflow.gcp.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn" ) def test_create_instance_when_not_exists(self, mock_get_conn, mock_project_id): mock_get_conn.return_value.get_instance.side_effect = [ NotFound("Instnace not found"), Instance(name=TEST_NAME), ] mock_get_conn.return_value.create_instance.return_value.result.return_value = Instance(name=TEST_NAME) result = self.hook.create_instance( location=TEST_LOCATION, instance_id=TEST_INSTANCE_ID, instance=Instance(name=TEST_NAME), retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) 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), ] ) mock_get_conn.return_value.create_instance.assert_called_once_with( instance=Instance( name=TEST_NAME, labels={"airflow-version": "v" + version.version.replace(".", "-").replace("+", "-")}, ), instance_id=TEST_INSTANCE_ID, metadata=TEST_METADATA, parent=TEST_PARENT_DEFAULT_PROJECT_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, ) self.assertEqual(Instance(name=TEST_NAME), result) @mock.patch( 'airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST ) @mock.patch( # type: ignore "airflow.gcp.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn" ) def test_delete_instance(self, mock_get_conn, mock_project_id): self.hook.delete_instance( location=TEST_LOCATION, instance=TEST_INSTANCE_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) mock_get_conn.return_value.delete_instance.assert_called_once_with( name=TEST_NAME_DEFAULT_PROJECT_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA ) @mock.patch( 'airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST ) @mock.patch( # type: ignore "airflow.gcp.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn" ) def test_get_instance(self, mock_get_conn, mock_project_id): self.hook.get_instance( location=TEST_LOCATION, instance=TEST_INSTANCE_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) mock_get_conn.return_value.get_instance.assert_called_once_with( name=TEST_NAME_DEFAULT_PROJECT_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA ) @mock.patch( 'airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST ) @mock.patch( # type: ignore "airflow.gcp.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn" ) def test_list_instances(self, mock_get_conn, mock_project_id): self.hook.list_instances( location=TEST_LOCATION, page_size=TEST_PAGE_SIZE, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) mock_get_conn.return_value.list_instances.assert_called_once_with( parent=TEST_PARENT_DEFAULT_PROJECT_ID, page_size=TEST_PAGE_SIZE, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) @mock.patch( 'airflow.gcp.hooks.base.GoogleCloudBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST ) @mock.patch( # type: ignore "airflow.gcp.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn" ) def test_update_instance(self, mock_get_conn, mock_project_id): self.hook.update_instance( update_mask=TEST_UPDATE_MASK, instance=Instance(name=TEST_NAME), location=TEST_LOCATION, instance_id=TEST_INSTANCE_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) mock_get_conn.return_value.update_instance.assert_called_once_with( update_mask=TEST_UPDATE_MASK, instance=Instance(name=TEST_NAME_DEFAULT_PROJECT_ID), retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, )