def create_cluster(self, project_id, region, cluster): self._check_region_matches_endpoint(region) # convert dict to object if not isinstance(cluster, Cluster): cluster = Cluster(**cluster) if cluster.project_id: if cluster.project_id != project_id: raise InvalidArgument( 'If provided, CreateClusterRequest.cluster.project_id must' ' match CreateClusterRequest.project_id') else: cluster.project_id = project_id if not cluster.cluster_name: raise InvalidArgument('Cluster name is required') # initialize cluster status cluster.status.state = _cluster_state_value('CREATING') cluster_key = (project_id, region, cluster.cluster_name) if cluster_key in self.mock_clusters: raise AlreadyExists('Already exists: Cluster ' + _cluster_path(*cluster_key)) self.mock_clusters[cluster_key] = cluster
def test_create_preexisting_topic_failifexists(self, mock_service): mock_service.return_value.create_topic.side_effect = AlreadyExists( 'Topic already exists: %s' % TEST_TOPIC ) with self.assertRaises(PubSubException) as e: self.pubsub_hook.create_topic(project_id=TEST_PROJECT, topic=TEST_TOPIC, fail_if_exists=True) self.assertEqual(str(e.exception), 'Topic already exists: %s' % TEST_TOPIC)
def test_execute_if_cluster_exists_in_error_state(self, mock_hook): mock_hook.return_value.create_cluster.side_effect = [AlreadyExists("test")] cluster_status = mock_hook.return_value.get_cluster.return_value.status cluster_status.state = 0 cluster_status.ERROR = 0 op = DataprocCreateClusterOperator( task_id=TASK_ID, region=GCP_LOCATION, project_id=GCP_PROJECT, cluster_config=CONFIG, labels=LABELS, cluster_name=CLUSTER_NAME, delete_on_error=True, gcp_conn_id=GCP_CONN_ID, retry=RETRY, timeout=TIMEOUT, metadata=METADATA, request_id=REQUEST_ID, ) with self.assertRaises(AirflowException): op.execute(context={}) mock_hook.return_value.diagnose_cluster.assert_called_once_with( region=GCP_LOCATION, project_id=GCP_PROJECT, cluster_name=CLUSTER_NAME ) mock_hook.return_value.delete_cluster.assert_called_once_with( region=GCP_LOCATION, project_id=GCP_PROJECT, cluster_name=CLUSTER_NAME )
def test_execute_if_cluster_exists_in_deleting_state( self, mock_hook, mock_get_cluster, mock_create_cluster, mock_generator ): cluster = mock.MagicMock() cluster.status.state = 0 cluster.status.DELETING = 0 cluster2 = mock.MagicMock() cluster2.status.state = 0 cluster2.status.ERROR = 0 mock_create_cluster.side_effect = [AlreadyExists("test"), cluster2] mock_generator.return_value = [0] mock_get_cluster.side_effect = [cluster, NotFound("test")] op = DataprocCreateClusterOperator( task_id=TASK_ID, region=GCP_LOCATION, project_id=GCP_PROJECT, cluster_config=CONFIG, labels=LABELS, cluster_name=CLUSTER_NAME, delete_on_error=True, gcp_conn_id=GCP_CONN_ID, ) with self.assertRaises(AirflowException): op.execute(context={}) calls = [mock.call(mock_hook.return_value), mock.call(mock_hook.return_value)] mock_get_cluster.assert_has_calls(calls) mock_create_cluster.assert_has_calls(calls) mock_hook.return_value.diagnose_cluster.assert_called_once_with( region=GCP_LOCATION, project_id=GCP_PROJECT, cluster_name=CLUSTER_NAME )
def test_execute_if_cluster_exists(self, mock_hook): mock_hook.return_value.create_cluster.side_effect = [ AlreadyExists("test") ] op = DataprocCreateClusterOperator( task_id=TASK_ID, region=GCP_LOCATION, project_id=GCP_PROJECT, cluster=CLUSTER, gcp_conn_id=GCP_CONN_ID, retry=RETRY, timeout=TIMEOUT, metadata=METADATA, request_id=REQUEST_ID, ) op.execute(context={}) mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID) mock_hook.return_value.create_cluster.assert_called_once_with( region=GCP_LOCATION, project_id=GCP_PROJECT, cluster=CLUSTER, request_id=REQUEST_ID, retry=RETRY, timeout=TIMEOUT, metadata=METADATA, ) mock_hook.return_value.get_cluster.assert_called_once_with( region=GCP_LOCATION, project_id=GCP_PROJECT, cluster_name=CLUSTER_NAME, retry=RETRY, timeout=TIMEOUT, metadata=METADATA, )
def test_create_subscription_nofailifexists(self, mock_service): mock_service.create_subscription.side_effect = AlreadyExists( 'Subscription already exists: %s' % EXPANDED_SUBSCRIPTION) response = self.pubsub_hook.create_subscription( project_id=TEST_PROJECT, topic=TEST_TOPIC, subscription=TEST_SUBSCRIPTION) self.assertEqual(TEST_SUBSCRIPTION, response)
def test_create_preexisting_topic_failifexists(self, mock_service): mock_service.return_value.create_topic.side_effect = AlreadyExists( f'Topic already exists: {TEST_TOPIC}') with pytest.raises(PubSubException) as ctx: self.pubsub_hook.create_topic(project_id=TEST_PROJECT, topic=TEST_TOPIC, fail_if_exists=True) assert str(ctx.value) == f'Topic already exists: {TEST_TOPIC}'
def test_create_subscription_failifexists(self, mock_service): mock_service.create_subscription.side_effect = AlreadyExists( 'Subscription already exists: %s' % EXPANDED_SUBSCRIPTION ) with self.assertRaises(PubSubException) as e: self.pubsub_hook.create_subscription( project_id=TEST_PROJECT, topic=TEST_TOPIC, subscription=TEST_SUBSCRIPTION, fail_if_exists=True ) self.assertEqual(str(e.exception), 'Subscription already exists: %s' % EXPANDED_SUBSCRIPTION)
def create_instance(self, parent, instance_id, instance, metadata=None): from google.api_core.exceptions import AlreadyExists, Unknown self._created_instance = (parent, instance_id, instance, metadata) if self._rpc_error: raise Unknown("error") if self._create_instance_conflict: raise AlreadyExists("conflict") return self._create_instance_response
class TestCloudVisionReferenceImageCreate(unittest.TestCase): @mock.patch( 'airflow.providers.google.cloud.operators.vision.CloudVisionHook', ) def test_minimal_green_path(self, mock_hook): mock_hook.return_value.create_reference_image.return_value = {} op = CloudVisionCreateReferenceImageOperator( location=LOCATION_TEST, product_id=PRODUCT_ID_TEST, reference_image=REFERENCE_IMAGE_TEST, task_id='id', ) op.execute(context=None) mock_hook.assert_called_once_with( gcp_conn_id=GCP_CONN_ID, impersonation_chain=None, ) mock_hook.return_value.create_reference_image.assert_called_once_with( location=LOCATION_TEST, product_id=PRODUCT_ID_TEST, reference_image=REFERENCE_IMAGE_TEST, reference_image_id=None, project_id=None, retry=None, timeout=None, metadata=None, ) @mock.patch( 'airflow.providers.google.cloud.operators.vision.CloudVisionHook', **{ 'return_value.create_reference_image.side_effect': AlreadyExists("MESSAGe") }, ) def test_already_exists(self, mock_hook): # Exception AlreadyExists not raised, caught in the operator's execute() - idempotence op = CloudVisionCreateReferenceImageOperator( location=LOCATION_TEST, product_id=PRODUCT_ID_TEST, reference_image=REFERENCE_IMAGE_TEST, task_id='id', ) op.execute(context=None) mock_hook.assert_called_once_with( gcp_conn_id=GCP_CONN_ID, impersonation_chain=None, ) mock_hook.return_value.create_reference_image.assert_called_once_with( location=LOCATION_TEST, product_id=PRODUCT_ID_TEST, reference_image=REFERENCE_IMAGE_TEST, reference_image_id=None, project_id=None, retry=None, timeout=None, metadata=None, )
def test_create_cluster_already_exists(self, wait_mock, convert_mock, log_mock): from google.api_core.exceptions import AlreadyExists # To force an error message = 'Already Exists' self.gke_hook._client.create_cluster.side_effect = AlreadyExists(message=message) self.gke_hook.create_cluster({}) wait_mock.assert_not_called() self.assertEqual(convert_mock.call_count, 1) log_mock.info.assert_any_call("Assuming Success: %s", message)
def test_create_subscription_failifexists(self, mock_service): mock_service.create_subscription.side_effect = AlreadyExists( f'Subscription already exists: {EXPANDED_SUBSCRIPTION}') with pytest.raises(PubSubException) as ctx: self.pubsub_hook.create_subscription( project_id=TEST_PROJECT, topic=TEST_TOPIC, subscription=TEST_SUBSCRIPTION, fail_if_exists=True) assert str(ctx.value ) == f'Subscription already exists: {EXPANDED_SUBSCRIPTION}'
def test_already_exists(self, mock_hook): mock_hook.return_value.create_product_set.side_effect = AlreadyExists(message='') # Exception AlreadyExists not raised, caught in the operator's execute() - idempotence op = CloudVisionCreateProductSetOperator( location=LOCATION_TEST, product_set=PRODUCTSET_TEST, product_set_id=PRODUCTSET_ID_TEST, project_id='mock-project-id', task_id='id', ) result = op.execute(None) self.assertEqual(PRODUCTSET_ID_TEST, result)
def test_already_exists(self, create_product_mock, get_conn): get_conn.return_value = {} create_product_mock.side_effect = AlreadyExists(message='') # Exception AlreadyExists not raised, caught in the operator's execute() - idempotence op = CloudVisionProductCreateOperator( location=LOCATION_TEST, product=PRODUCT_TEST, product_id=PRODUCT_ID_TEST, project_id='mock-project-id', task_id='id', ) result = op.execute(None) self.assertEqual(PRODUCT_ID_TEST, result)
def create_database(self, parent, create_statement, extra_statements=None, metadata=None): from google.api_core.exceptions import AlreadyExists, NotFound, Unknown self._created_database = ( parent, create_statement, extra_statements, metadata) if self._rpc_error: raise Unknown('error') if self._create_database_conflict: raise AlreadyExists('conflict') if self._database_not_found: raise NotFound('not found') return self._create_database_response
def test_create_cluster_already_exists(self, wait_mock, convert_mock, log_mock, mock_get_credentials): from google.api_core.exceptions import AlreadyExists # To force an error message = 'Already Exists' self.gke_hook._client.create_cluster.side_effect = AlreadyExists( message=message) self.gke_hook.create_cluster(cluster={}, project_id=TEST_GCP_PROJECT_ID) wait_mock.assert_not_called() assert convert_mock.call_count == 1 log_mock.info.assert_any_call("Assuming Success: %s", message)
def create_cluster(self, project_id, region, cluster): self._check_region_matches_endpoint(region) # convert dict to object if not isinstance(cluster, Cluster): cluster = Cluster(**cluster) if cluster.project_id: if cluster.project_id != project_id: raise InvalidArgument( 'If provided, CreateClusterRequest.cluster.project_id must' ' match CreateClusterRequest.project_id') else: cluster.project_id = project_id if not cluster.cluster_name: raise InvalidArgument('Cluster name is required') # add in default disk config for x in ('master', 'worker', 'secondary_worker'): field = x + '_config' conf = getattr(cluster.config, field, None) if conf and str(conf): # empty DiskConfigs are still true-ish if not conf.disk_config: conf.disk_config = DiskConfig() if not conf.disk_config.boot_disk_size_gb: conf.disk_config.boot_disk_size_gb = _DEFAULT_DISK_SIZE_GB # update gce_cluster_config gce_config = cluster.config.gce_cluster_config # add in default scopes and sort scopes = set(gce_config.service_account_scopes) if not scopes: scopes.update(_DEFAULT_SCOPES) scopes.update(_MANDATORY_SCOPES) gce_config.service_account_scopes[:] = sorted(scopes) # initialize cluster status cluster.status.state = _cluster_state_value('CREATING') cluster_key = (project_id, region, cluster.cluster_name) if cluster_key in self.mock_clusters: raise AlreadyExists('Already exists: Cluster ' + _cluster_path(*cluster_key)) self.mock_clusters[cluster_key] = cluster
def create_instance(self, create_in_production, cluster_name, location_id, nr_nodes, use_ssd_storage, timeout=100): """To create an instance, you also have to configure cluster parameters. If the cluster already exists, an AlreadyExists exception is raised. Example configuration: - create_in_production=False - cluster_name='test_cluster' - location_id='us-central1-f' - nr_nodes=1 - ssd_storage=False Args: create_in_production (bool): If True, then production instance is created. If False, a development instance will be created; note that nr_nodes will not be used as a parameter in that case. cluster_name (str): name of the cluster that will be created within the instance. location_id (str): the zone to create the cluster in. nr_nodes (int): only meaningful for production instance. Otherwise, in development instance, will be set to 1. use_ssd_storage (bool): if True, then SSD storage is used, otherwise HDD storage. timeout (int): Returns: None """ if self.instance.exists(): raise AlreadyExists("Instance '{}' already exists.".format( self.instance_id)) logging.info("Creating instance '%s'.", self.instance_id) # instance configurations production, nr_nodes = (bt_enums.Instance.Type.PRODUCTION, nr_nodes) if create_in_production else ( bt_enums.Instance.Type.DEVELOPMENT, None) self.instance = self.client.instance(instance_id=self.instance_id, display_name=self.instance_id, instance_type=production) # cluster configurations cluster = self.create_cluster_config(cluster_name, location_id, nr_nodes, use_ssd_storage) # Create the instance with a cluster operation = self.instance.create(clusters=[cluster]) # We want to make sure the operation completes. operation.result(timeout=timeout)
def test_create_topic_flow_when_topic_exists(self, mock_json_renderer, mock_google_client): with self.settings(GCLOUD_PUBSUB_PROJECT_ID="pubsub"): mock_publisher = mock_google_client() mock_publisher.topic_path.return_value = "topic_path" mock_json_renderer() mock_publisher.create_topic.side_effect = AlreadyExists( "Topic Exists!") with self.assertLogs( "google_pubsub_adapter.publisher.google_pubsub_publisher", level="INFO") as mock_log: publisher = GooglePubsubPublisher() publisher.create_topic("topic") self.assertIn("already exists", mock_log.output[0])
def test_execute_if_cluster_exists_do_not_use(self, mock_hook): mock_hook.return_value.create_cluster.side_effect = [AlreadyExists("test")] mock_hook.return_value.get_cluster.return_value.status.state = 0 op = DataprocCreateClusterOperator( task_id=TASK_ID, region=GCP_LOCATION, project_id=GCP_PROJECT, cluster=CLUSTER, gcp_conn_id=GCP_CONN_ID, retry=RETRY, timeout=TIMEOUT, metadata=METADATA, request_id=REQUEST_ID, use_if_exists=False, ) with self.assertRaises(AlreadyExists): op.execute(context={})
def test_execute_if_cluster_exists(self, mock_hook, to_dict_mock): mock_hook.return_value.create_cluster.side_effect = [ AlreadyExists("test") ] mock_hook.return_value.get_cluster.return_value.status.state = 0 op = DataprocCreateClusterOperator( task_id=TASK_ID, region=GCP_LOCATION, project_id=GCP_PROJECT, cluster_config=CONFIG, labels=LABELS, cluster_name=CLUSTER_NAME, gcp_conn_id=GCP_CONN_ID, retry=RETRY, timeout=TIMEOUT, metadata=METADATA, request_id=REQUEST_ID, impersonation_chain=IMPERSONATION_CHAIN, ) op.execute(context={}) mock_hook.assert_called_once_with( gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) mock_hook.return_value.create_cluster.assert_called_once_with( region=GCP_LOCATION, project_id=GCP_PROJECT, cluster_config=CONFIG, labels=LABELS, cluster_name=CLUSTER_NAME, request_id=REQUEST_ID, retry=RETRY, timeout=TIMEOUT, metadata=METADATA, ) mock_hook.return_value.get_cluster.assert_called_once_with( region=GCP_LOCATION, project_id=GCP_PROJECT, cluster_name=CLUSTER_NAME, retry=RETRY, timeout=TIMEOUT, metadata=METADATA, ) to_dict_mock.assert_called_once_with( mock_hook.return_value.get_cluster.return_value)
def submit_job(self, project_id, region, job): self._check_region_matches_endpoint(region) # convert dict to object if not isinstance(job, Job): job = Job(**job) if not (project_id and job.reference.job_id): raise NotImplementedError('generation of job IDs not implemented') job_id = job.reference.job_id if not job.placement.cluster_name: raise InvalidArgument('Cluster name is required') # cluster must exist cluster_key = (project_id, region, job.placement.cluster_name) if cluster_key not in self.mock_clusters: raise NotFound('Not Found: Cluster ' + _cluster_path(*cluster_key)) if not job.hadoop_job: raise NotImplementedError('only hadoop jobs are supported') if job.reference.project_id: if job.reference.project_id != project_id: raise InvalidArgument( 'If provided, SubmitJobRequest.job.job_reference' '.project_id must match SubmitJobRequest.project_id') else: job.reference.project_id = project_id job.status.state = _job_state_value('SETUP_DONE') job_key = (project_id, region, job_id) if job_key in self.mock_jobs: raise AlreadyExists( 'Already exists: Job ' + _job_path(*job_key)) self.mock_jobs[job_key] = job return deepcopy(job)
class TestCatchHttpException(unittest.TestCase): # pylint:disable=no-method-argument,unused-argument @parameterized.expand([ ("no_exception", None, LoggingMixin, None, None), ("raise_airflowexception", MovedPermanently("MESSAGE"), LoggingMixin, None, AirflowException), ("raise_airflowexception", RetryError("MESSAGE", cause=Exception("MESSAGE")), LoggingMixin, None, AirflowException), ("raise_airflowexception", ValueError("MESSAGE"), LoggingMixin, None, AirflowException), ("raise_alreadyexists", AlreadyExists("MESSAGE"), LoggingMixin, None, AlreadyExists), ("raise_http_error", HttpError(mock.Mock(**{"reason.return_value": None}), b"CONTENT"), BaseHook, { "source": None }, AirflowException), ]) def test_catch_exception(self, name, exception, base_class, base_class_args, assert_raised): self.called = False # pylint:disable=attribute-defined-outside-init class FixtureClass(base_class): @hook.GoogleCloudBaseHook.catch_http_exception def test_fixture(*args, **kwargs): # pylint:disable=unused-argument,no-method-argument self.called = True # pylint:disable=attribute-defined-outside-init if exception is not None: raise exception if assert_raised is None: FixtureClass(base_class_args).test_fixture() else: with self.assertRaises(assert_raised): FixtureClass(base_class_args).test_fixture() self.assertTrue(self.called)
def create_cluster(self, project_id, region, cluster): self._check_region_matches_endpoint(region) # convert dict to object if not isinstance(cluster, Cluster): cluster = Cluster(**cluster) if cluster.project_id: if cluster.project_id != project_id: raise InvalidArgument( 'If provided, CreateClusterRequest.cluster.project_id must' ' match CreateClusterRequest.project_id') else: cluster.project_id = project_id if not cluster.cluster_name: raise InvalidArgument('Cluster name is required') # add in default disk config for x in ('master', 'worker', 'secondary_worker'): field = x + '_config' conf = getattr(cluster.config, field, None) if conf and str(conf): # empty DiskConfigs are still true-ish if not conf.disk_config: conf.disk_config = DiskConfig() if not conf.disk_config.boot_disk_size_gb: conf.disk_config.boot_disk_size_gb = _DEFAULT_DISK_SIZE_GB # update gce_cluster_config gce_config = cluster.config.gce_cluster_config # check region and zone_uri if region == 'global': if gce_config.zone_uri: cluster_region = _zone_to_region(gce_config.zone_uri) else: raise InvalidArgument( "Must specify a zone in GCE configuration" " when using 'regions/global'") else: cluster_region = region # add in default scopes and sort scopes = set(gce_config.service_account_scopes) if not scopes: scopes.update(_DEFAULT_SCOPES) scopes.update(_MANDATORY_SCOPES) gce_config.service_account_scopes[:] = sorted(scopes) # handle network_uri and subnetwork_uri if gce_config.network_uri and gce_config.subnetwork_uri: raise InvalidArgument('GceClusterConfiguration cannot contain both' ' Network URI and Subnetwork URI') if not (gce_config.network_uri or gce_config.subnetwork_uri): gce_config.network_uri = 'default' if gce_config.network_uri: gce_config.network_uri = _fully_qualify_network_uri( gce_config.network_uri, project_id) if gce_config.subnetwork_uri: gce_config.subnetwork_uri = _fully_qualify_subnetwork_uri( gce_config.subnetwork_uri, project_id, region) # add in default cluster properties props = cluster.config.software_config.properties for k, v in _DEFAULT_CLUSTER_PROPERTIES.items(): if k not in props: props[k] = v # initialize cluster status cluster.status.state = _cluster_state_value('CREATING') cluster_key = (project_id, region, cluster.cluster_name) if cluster_key in self.mock_clusters: raise AlreadyExists('Already exists: Cluster ' + _cluster_path(*cluster_key)) self.mock_clusters[cluster_key] = cluster
def test_fixutre(*args, **kwargs): self.called = True raise AlreadyExists("MESSAGE")
def upload(self, metadata_file_name): # type: (str) -> None metadata_file_name = os.path.abspath(metadata_file_name) metadata_directory = os.path.dirname(metadata_file_name) with open(metadata_file_name, 'r') as file: metadata = json.load(file) if not isinstance(metadata, dict): raise RuntimeError('Metadata is not a dictionary in file %s' % metadata_file_name) structured_to_send = {} uploads_to_send = {} nb_no_category = 0 nb_sendings = 0 for image_path, image_metadata in sorted(metadata.items()): image_basename = os.path.basename(image_path) sending = Sending() sending.category = image_metadata.get('category', None) sending.location = image_metadata['location'] sending.timestamp = image_metadata['timestamp'] sending.firebase_path = image_basename.replace('_', '/') sending.collection_id, sending.upload_id, image_name = sending.firebase_path.split( '/') sending.image_id = int(os.path.splitext(image_name)[0]) sending.local_path = os.path.join(metadata_directory, image_basename) assert sending.timestamp, 'Got an invalid timestamp' nb_no_category += sending.category is None nb_sendings += 1 structured_to_send.setdefault(sending.collection_id, {}).setdefault( sending.upload_id, []).append(sending) # Check categories. if nb_no_category == nb_sendings: inferred_category = os.path.basename(metadata_directory) print('Getting category from metadata containing folder', inferred_category) for upload_id_to_sendings in structured_to_send.values(): for sendings in upload_id_to_sendings.values(): for sending in sendings: sending.category = inferred_category elif nb_no_category != 0: raise RuntimeError( 'Invalid metadata: all images should either have a category or ' 'no category specified (to be retrieved from metadata folder name).' ) # Check timestamps. for collection_id, uploads in structured_to_send.items(): for upload_id, images in uploads.items(): timestamps = {sending.timestamp for sending in images} if len(timestamps) != 1: raise RuntimeError( 'No same timestamp for all images in %s' % upload_id) timestamp = images[0].timestamp # Update timestamp to current date if timestamp is unknown if timestamp == ImageInfo.UNKNOWN_CATEGORY: timestamp = datetime.now().isoformat() uploads_to_send.setdefault(collection_id, []).append( UploadToSend(upload_id, timestamp, images)) for collection_id, uploads in uploads_to_send.items(): col = self.__database.collection(collection_id) for upload in sorted(uploads, key=lambda u: u.upload_id): doc = col.document(upload.upload_id) doc_just_created = False try: doc.create({ 'timestamp': datetime.fromisoformat(upload.timestamp), 'images': None }) doc_just_created = True for sending in upload.images: # type: Sending blob = self.__storage.blob(sending.firebase_path) if blob.exists(): raise AlreadyExists('An image already exists: %s' % sending.firebase_path) for sending in sorted( upload.images, key=lambda s: s.image_id): # type: Sending blob = self.__storage.blob(sending.firebase_path) blob.upload_from_filename(sending.local_path) sending.url = blob.public_url print('UPLOADED', sending.firebase_path) doc.update({ 'images': [ sending.to_upload() for sending in sorted(upload.images, key=lambda s: s.image_id) ] }) print('CREATED UPLOAD', upload.upload_id) except AlreadyExists as exc: print('CANNOT SEND UPLOAD', upload.upload_id, exc) if doc_just_created: doc.delete()
class TestCloudDataCatalogCreateEntryOperator(TestCase): @mock.patch( "airflow.providers.google.cloud.operators.datacatalog.CloudDataCatalogHook", **{"return_value.create_entry.return_value": TEST_ENTRY}, ) def test_assert_valid_hook_call(self, mock_hook) -> None: task = CloudDataCatalogCreateEntryOperator( task_id="task_id", location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry_id=TEST_ENTRY_ID, entry=TEST_ENTRY, project_id=TEST_PROJECT_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, gcp_conn_id=TEST_GCP_CONN_ID, impersonation_chain=TEST_IMPERSONATION_CHAIN, ) ti = mock.MagicMock() result = task.execute(context={"task_instance": ti}) mock_hook.assert_called_once_with( gcp_conn_id=TEST_GCP_CONN_ID, impersonation_chain=TEST_IMPERSONATION_CHAIN, ) mock_hook.return_value.create_entry.assert_called_once_with( location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry_id=TEST_ENTRY_ID, entry=TEST_ENTRY, project_id=TEST_PROJECT_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) ti.xcom_push.assert_called_once_with(key="entry_id", value=TEST_ENTRY_ID) self.assertEqual(TEST_ENTRY_DICT, result) @mock.patch( "airflow.providers.google.cloud.operators.datacatalog.CloudDataCatalogHook", **{ "return_value.create_entry.side_effect": AlreadyExists(message="message"), "return_value.get_entry.return_value": TEST_ENTRY, }, ) def test_assert_valid_hook_call_when_exists(self, mock_hook) -> None: task = CloudDataCatalogCreateEntryOperator( task_id="task_id", location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry_id=TEST_ENTRY_ID, entry=TEST_ENTRY, project_id=TEST_PROJECT_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, gcp_conn_id=TEST_GCP_CONN_ID, impersonation_chain=TEST_IMPERSONATION_CHAIN, ) ti = mock.MagicMock() result = task.execute(context={"task_instance": ti}) mock_hook.assert_called_once_with( gcp_conn_id=TEST_GCP_CONN_ID, impersonation_chain=TEST_IMPERSONATION_CHAIN, ) mock_hook.return_value.create_entry.assert_called_once_with( location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry_id=TEST_ENTRY_ID, entry=TEST_ENTRY, project_id=TEST_PROJECT_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) mock_hook.return_value.get_entry.assert_called_once_with( location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, project_id=TEST_PROJECT_ID, retry=TEST_RETRY, timeout=TEST_TIMEOUT, metadata=TEST_METADATA, ) ti.xcom_push.assert_called_once_with(key="entry_id", value=TEST_ENTRY_ID) self.assertEqual(TEST_ENTRY_DICT, result)
def test_create_preexisting_topic_nofailifexists(self, mock_service): mock_service.return_value.create_topic.side_effect = AlreadyExists( 'Topic already exists: %s' % EXPANDED_TOPIC) self.pubsub_hook.create_topic(project_id=TEST_PROJECT, topic=TEST_TOPIC)