def execute(self, context) -> None: hook = BigtableHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) instance = hook.get_instance(project_id=self.project_id, instance_id=self.instance_id) if not instance: raise AirflowException( "Dependency: instance '{}' does not exist in project '{}'.".format( self.instance_id, self.project_id ) ) try: hook.create_table( instance=instance, table_id=self.table_id, initial_split_keys=self.initial_split_keys, column_families=self.column_families, ) except google.api_core.exceptions.AlreadyExists: if not self._compare_column_families(hook, instance): raise AirflowException( f"Table '{self.table_id}' already exists with different Column Families." ) self.log.info("The table '%s' already exists. Consider it as created", self.table_id)
def execute(self, context): hook = BigtableHook(gcp_conn_id=self.gcp_conn_id) instance = hook.get_instance(project_id=self.project_id, instance_id=self.instance_id) if instance: # Based on Instance.__eq__ instance with the same ID and client is # considered as equal. self.log.info( "The instance '%s' already exists in this project. " "Consider it as created", self.instance_id) return try: hook.create_instance( project_id=self.project_id, instance_id=self.instance_id, main_cluster_id=self.main_cluster_id, main_cluster_zone=self.main_cluster_zone, replica_cluster_id=self.replica_cluster_id, replica_cluster_zone=self.replica_cluster_zone, instance_display_name=self.instance_display_name, instance_type=self.instance_type, instance_labels=self.instance_labels, cluster_nodes=self.cluster_nodes, cluster_storage_type=self.cluster_storage_type, timeout=self.timeout, ) except google.api_core.exceptions.GoogleAPICallError as e: self.log.error('An error occurred. Exiting.') raise e
def execute(self, context) -> None: hook = BigtableHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) instance = hook.get_instance(project_id=self.project_id, instance_id=self.instance_id) if not instance: raise AirflowException( f"Dependency: instance '{self.instance_id}' does not exist.") try: hook.delete_table( project_id=self.project_id, instance_id=self.instance_id, table_id=self.table_id, ) except google.api_core.exceptions.NotFound: # It's OK if table doesn't exists. self.log.info( "The table '%s' no longer exists. Consider it as deleted", self.table_id) except google.api_core.exceptions.GoogleAPICallError as e: self.log.error('An error occurred. Exiting.') raise e
def poke(self, context: dict) -> bool: hook = BigtableHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) instance = hook.get_instance(project_id=self.project_id, instance_id=self.instance_id) if not instance: self.log.info("Dependency: instance '%s' does not exist.", self.instance_id) return False try: cluster_states = hook.get_cluster_states_for_table(instance=instance, table_id=self.table_id) except google.api_core.exceptions.NotFound: self.log.info( "Dependency: table '%s' does not exist in instance '%s'.", self.table_id, self.instance_id ) return False ready_state = ClusterState(enums.Table.ClusterState.ReplicationState.READY) is_table_replicated = True for cluster_id in cluster_states.keys(): if cluster_states[cluster_id] != ready_state: self.log.info("Table '%s' is not yet replicated on cluster '%s'.", self.table_id, cluster_id) is_table_replicated = False if not is_table_replicated: return False self.log.info("Table '%s' is replicated.", self.table_id) return True
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.bigtable_hook_no_default_project_id = BigtableHook( gcp_conn_id='test')
def execute(self, context): hook = BigtableHook(gcp_conn_id=self.gcp_conn_id) try: hook.delete_instance(project_id=self.project_id, instance_id=self.instance_id) except google.api_core.exceptions.NotFound: self.log.info( "The instance '%s' does not exist in project '%s'. " "Consider it as deleted", self.instance_id, self.project_id) except google.api_core.exceptions.GoogleAPICallError as e: self.log.error('An error occurred. Exiting.') raise e
def execute(self, context): hook = BigtableHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain,) instance = hook.get_instance(project_id=self.project_id, instance_id=self.instance_id) if not instance: raise AirflowException("Dependency: instance '{}' does not exist.".format(self.instance_id)) try: hook.update_cluster(instance=instance, cluster_id=self.cluster_id, nodes=self.nodes) except google.api_core.exceptions.NotFound: raise AirflowException( "Dependency: cluster '{}' does not exist for instance '{}'.".format( self.cluster_id, self.instance_id ) ) except google.api_core.exceptions.GoogleAPICallError as e: self.log.error('An error occurred. Exiting.') raise e
def execute(self, context): hook = BigtableHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain,) instance = hook.get_instance(project_id=self.project_id, instance_id=self.instance_id) if not instance: raise AirflowException(f"Dependency: instance '{self.instance_id}' does not exist.") try: hook.update_instance( project_id=self.project_id, instance_id=self.instance_id, instance_display_name=self.instance_display_name, instance_type=self.instance_type, instance_labels=self.instance_labels, timeout=self.timeout, ) except google.api_core.exceptions.GoogleAPICallError as e: self.log.error('An error occurred. Exiting.') raise e
class TestBigtableHookNoDefaultProjectId(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.bigtable_hook_no_default_project_id = BigtableHook(gcp_conn_id='test') @mock.patch( "airflow.providers.google.cloud.hooks.bigtable.BigtableHook.client_info", new_callable=mock.PropertyMock, ) @mock.patch("airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_credentials") @mock.patch("airflow.providers.google.cloud.hooks.bigtable.Client") def test_bigtable_client_creation(self, mock_client, mock_get_creds, mock_client_info): result = self.bigtable_hook_no_default_project_id._get_client(GCP_PROJECT_ID_HOOK_UNIT_TEST) mock_client.assert_called_once_with( project=GCP_PROJECT_ID_HOOK_UNIT_TEST, credentials=mock_get_creds.return_value, client_info=mock_client_info.return_value, admin=True, ) self.assertEqual(mock_client.return_value, result) self.assertEqual(self.bigtable_hook_no_default_project_id._client, result) @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_get_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True res = self.bigtable_hook_no_default_project_id.get_instance( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=CBT_INSTANCE ) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() get_client.assert_called_once_with(project_id='example-project') self.assertIsNotNone(res) @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_delete_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True delete_method = instance_method.return_value.delete res = self.bigtable_hook_no_default_project_id.delete_instance( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=CBT_INSTANCE ) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() delete_method.assert_called_once_with() get_client.assert_called_once_with(project_id='example-project') self.assertIsNone(res) @mock.patch('google.cloud.bigtable.instance.Instance.create') @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_create_instance_overridden_project_id(self, get_client, instance_create): operation = mock.Mock() operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client) instance_create.return_value = operation res = self.bigtable_hook_no_default_project_id.create_instance( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=CBT_INSTANCE, main_cluster_id=CBT_CLUSTER, main_cluster_zone=CBT_ZONE, ) get_client.assert_called_once_with(project_id='example-project') instance_create.assert_called_once_with(clusters=mock.ANY) self.assertEqual(res.instance_id, 'instance') @mock.patch('google.cloud.bigtable.instance.Instance.update') @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_update_instance_overridden_project_id(self, get_client, instance_update): operation = mock.Mock() operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client) instance_update.return_value = operation res = self.bigtable_hook_no_default_project_id.update_instance( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=CBT_INSTANCE, instance_display_name=CBT_INSTANCE_DISPLAY_NAME, instance_type=CBT_INSTANCE_TYPE, instance_labels=CBT_INSTANCE_LABELS, ) get_client.assert_called_once_with(project_id='example-project') instance_update.assert_called_once_with() self.assertEqual(res.instance_id, 'instance') @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_delete_table_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists table_delete_method = instance_method.return_value.table.return_value.delete instance_exists_method.return_value = True self.bigtable_hook_no_default_project_id.delete_table( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=CBT_INSTANCE, table_id=CBT_TABLE ) get_client.assert_called_once_with(project_id='example-project') instance_exists_method.assert_called_once_with() table_delete_method.assert_called_once_with()
class TestBigtableHookDefaultProjectId(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.bigtable_hook_default_project_id = BigtableHook(gcp_conn_id='test') @mock.patch( "airflow.providers.google.cloud.hooks.bigtable.BigtableHook.client_info", new_callable=mock.PropertyMock, ) @mock.patch("airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_credentials") @mock.patch("airflow.providers.google.cloud.hooks.bigtable.Client") def test_bigtable_client_creation(self, mock_client, mock_get_creds, mock_client_info): result = self.bigtable_hook_default_project_id._get_client(GCP_PROJECT_ID_HOOK_UNIT_TEST) mock_client.assert_called_once_with( project=GCP_PROJECT_ID_HOOK_UNIT_TEST, credentials=mock_get_creds.return_value, client_info=mock_client_info.return_value, admin=True, ) self.assertEqual(mock_client.return_value, result) self.assertEqual(self.bigtable_hook_default_project_id._client, 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.bigtable.BigtableHook._get_client') def test_get_instance(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True res = self.bigtable_hook_default_project_id.get_instance( instance_id=CBT_INSTANCE, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() get_client.assert_called_once_with(project_id='example-project') self.assertIsNotNone(res) @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_get_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True res = self.bigtable_hook_default_project_id.get_instance( project_id='new-project', instance_id=CBT_INSTANCE ) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() get_client.assert_called_once_with(project_id='new-project') self.assertIsNotNone(res) @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.bigtable.BigtableHook._get_client') def test_get_instance_no_instance(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = False res = self.bigtable_hook_default_project_id.get_instance( instance_id=CBT_INSTANCE, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() get_client.assert_called_once_with(project_id='example-project') self.assertIsNone(res) @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.bigtable.BigtableHook._get_client') def test_delete_instance(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True delete_method = instance_method.return_value.delete res = self.bigtable_hook_default_project_id.delete_instance( instance_id=CBT_INSTANCE, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() delete_method.assert_called_once_with() get_client.assert_called_once_with(project_id='example-project') self.assertIsNone(res) @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_delete_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True delete_method = instance_method.return_value.delete res = self.bigtable_hook_default_project_id.delete_instance( project_id='new-project', instance_id=CBT_INSTANCE ) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() delete_method.assert_called_once_with() get_client.assert_called_once_with(project_id='new-project') self.assertIsNone(res) @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.bigtable.BigtableHook._get_client') def test_delete_instance_no_instance(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = False delete_method = instance_method.return_value.delete self.bigtable_hook_default_project_id.delete_instance( instance_id=CBT_INSTANCE, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() delete_method.assert_not_called() get_client.assert_called_once_with(project_id='example-project') @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('google.cloud.bigtable.instance.Instance.create') @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_create_instance(self, get_client, instance_create, mock_project_id): operation = mock.Mock() operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client) instance_create.return_value = operation res = self.bigtable_hook_default_project_id.create_instance( instance_id=CBT_INSTANCE, main_cluster_id=CBT_CLUSTER, main_cluster_zone=CBT_ZONE, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) get_client.assert_called_once_with(project_id='example-project') instance_create.assert_called_once_with(clusters=mock.ANY) self.assertEqual(res.instance_id, 'instance') @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('google.cloud.bigtable.instance.Instance.cluster') @mock.patch('google.cloud.bigtable.instance.Instance.create') @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_create_instance_with_one_replica_cluster( self, get_client, instance_create, cluster, mock_project_id ): operation = mock.Mock() operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client) instance_create.return_value = operation res = self.bigtable_hook_default_project_id.create_instance( instance_id=CBT_INSTANCE, main_cluster_id=CBT_CLUSTER, main_cluster_zone=CBT_ZONE, replica_cluster_id=CBT_REPLICA_CLUSTER_ID, replica_cluster_zone=CBT_REPLICA_CLUSTER_ZONE, cluster_nodes=1, cluster_storage_type=enums.StorageType.SSD, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) cluster.assert_has_calls( [ unittest.mock.call(CBT_CLUSTER, CBT_ZONE, 1, enums.StorageType.SSD), unittest.mock.call( CBT_REPLICA_CLUSTER_ID, CBT_REPLICA_CLUSTER_ZONE, 1, enums.StorageType.SSD ), ], any_order=True, ) get_client.assert_called_once_with(project_id='example-project') instance_create.assert_called_once_with(clusters=mock.ANY) self.assertEqual(res.instance_id, 'instance') @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('google.cloud.bigtable.instance.Instance.cluster') @mock.patch('google.cloud.bigtable.instance.Instance.create') @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_create_instance_with_multiple_replica_clusters( self, get_client, instance_create, cluster, mock_project_id ): operation = mock.Mock() operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client) instance_create.return_value = operation res = self.bigtable_hook_default_project_id.create_instance( instance_id=CBT_INSTANCE, main_cluster_id=CBT_CLUSTER, main_cluster_zone=CBT_ZONE, replica_clusters=CBT_REPLICATE_CLUSTERS, cluster_nodes=1, cluster_storage_type=enums.StorageType.SSD, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) cluster.assert_has_calls( [ unittest.mock.call(CBT_CLUSTER, CBT_ZONE, 1, enums.StorageType.SSD), unittest.mock.call('replica-1', 'us-west1-a', 1, enums.StorageType.SSD), unittest.mock.call('replica-2', 'us-central1-f', 1, enums.StorageType.SSD), unittest.mock.call('replica-3', 'us-east1-d', 1, enums.StorageType.SSD), ], any_order=True, ) get_client.assert_called_once_with(project_id='example-project') instance_create.assert_called_once_with(clusters=mock.ANY) self.assertEqual(res.instance_id, 'instance') @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('google.cloud.bigtable.instance.Instance.update') @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_update_instance(self, get_client, instance_update, mock_project_id): operation = mock.Mock() operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client) instance_update.return_value = operation res = self.bigtable_hook_default_project_id.update_instance( instance_id=CBT_INSTANCE, instance_display_name=CBT_INSTANCE_DISPLAY_NAME, instance_type=CBT_INSTANCE_TYPE, instance_labels=CBT_INSTANCE_LABELS, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) get_client.assert_called_once_with(project_id='example-project') instance_update.assert_called_once_with() self.assertEqual(res.instance_id, 'instance') @mock.patch('google.cloud.bigtable.instance.Instance.create') @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_create_instance_overridden_project_id(self, get_client, instance_create): operation = mock.Mock() operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client) instance_create.return_value = operation res = self.bigtable_hook_default_project_id.create_instance( project_id='new-project', instance_id=CBT_INSTANCE, main_cluster_id=CBT_CLUSTER, main_cluster_zone=CBT_ZONE, ) get_client.assert_called_once_with(project_id='new-project') instance_create.assert_called_once_with(clusters=mock.ANY) self.assertEqual(res.instance_id, 'instance') @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.bigtable.BigtableHook._get_client') def test_delete_table(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists table_delete_method = instance_method.return_value.table.return_value.delete instance_exists_method.return_value = True self.bigtable_hook_default_project_id.delete_table( instance_id=CBT_INSTANCE, table_id=CBT_TABLE, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) get_client.assert_called_once_with(project_id='example-project') instance_exists_method.assert_called_once_with() table_delete_method.assert_called_once_with() @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_delete_table_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists table_delete_method = instance_method.return_value.table.return_value.delete instance_exists_method.return_value = True self.bigtable_hook_default_project_id.delete_table( project_id='new-project', instance_id=CBT_INSTANCE, table_id=CBT_TABLE ) get_client.assert_called_once_with(project_id='new-project') instance_exists_method.assert_called_once_with() table_delete_method.assert_called_once_with() @mock.patch('google.cloud.bigtable.table.Table.create') @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_create_table(self, get_client, create): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True client = mock.Mock(Client) instance = google.cloud.bigtable.instance.Instance(instance_id=CBT_INSTANCE, client=client) self.bigtable_hook_default_project_id.create_table(instance=instance, table_id=CBT_TABLE) get_client.assert_not_called() create.assert_called_once_with([], {}) @mock.patch('google.cloud.bigtable.cluster.Cluster.update') @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_update_cluster(self, get_client, update): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True client = mock.Mock(Client) instance = google.cloud.bigtable.instance.Instance(instance_id=CBT_INSTANCE, client=client) self.bigtable_hook_default_project_id.update_cluster( instance=instance, cluster_id=CBT_CLUSTER, nodes=4 ) get_client.assert_not_called() update.assert_called_once_with() @mock.patch('google.cloud.bigtable.table.Table.list_column_families') @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_list_column_families(self, get_client, list_column_families): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True client = mock.Mock(Client) get_client.return_value = client instance = google.cloud.bigtable.instance.Instance(instance_id=CBT_INSTANCE, client=client) self.bigtable_hook_default_project_id.get_column_families_for_table( instance=instance, table_id=CBT_TABLE ) get_client.assert_not_called() list_column_families.assert_called_once_with() @mock.patch('google.cloud.bigtable.table.Table.get_cluster_states') @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_get_cluster_states(self, get_client, get_cluster_states): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True client = mock.Mock(Client) instance = google.cloud.bigtable.instance.Instance(instance_id=CBT_INSTANCE, client=client) self.bigtable_hook_default_project_id.get_cluster_states_for_table( instance=instance, table_id=CBT_TABLE ) get_client.assert_not_called() get_cluster_states.assert_called_once_with()
class TestBigtableHookNoDefaultProjectId(unittest.TestCase): def setUp(self): with mock.patch( 'airflow.providers.google.cloud.hooks.base.CloudBaseHook.__init__', new=mock_base_gcp_hook_no_default_project_id): self.bigtable_hook_no_default_project_id = BigtableHook( gcp_conn_id='test') @mock.patch( "airflow.providers.google.cloud.hooks.bigtable.BigtableHook.client_info", new_callable=mock.PropertyMock) @mock.patch( "airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_credentials" ) @mock.patch("airflow.providers.google.cloud.hooks.bigtable.Client") def test_bigtable_client_creation(self, mock_client, mock_get_creds, mock_client_info): result = self.bigtable_hook_no_default_project_id._get_client( GCP_PROJECT_ID_HOOK_UNIT_TEST) mock_client.assert_called_once_with( project=GCP_PROJECT_ID_HOOK_UNIT_TEST, credentials=mock_get_creds.return_value, client_info=mock_client_info.return_value, admin=True) self.assertEqual(mock_client.return_value, result) self.assertEqual(self.bigtable_hook_no_default_project_id._client, result) @mock.patch( 'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_get_instance_missing_project_id(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True with self.assertRaises(AirflowException) as cm: self.bigtable_hook_no_default_project_id.get_instance( instance_id=CBT_INSTANCE) instance_exists_method.assert_not_called() instance_method.assert_not_called() err = cm.exception self.assertIn("The project id must be passed", str(err)) @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_get_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True res = self.bigtable_hook_no_default_project_id.get_instance( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=CBT_INSTANCE) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() get_client.assert_called_once_with(project_id='example-project') self.assertIsNotNone(res) @mock.patch( 'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_delete_instance_missing_project_id(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists delete_method = instance_method.return_value.delete instance_exists_method.return_value = True with self.assertRaises(AirflowException) as cm: self.bigtable_hook_no_default_project_id.delete_instance( instance_id=CBT_INSTANCE) instance_exists_method.assert_not_called() instance_method.assert_not_called() delete_method.assert_not_called() err = cm.exception self.assertIn("The project id must be passed", str(err)) @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_delete_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True delete_method = instance_method.return_value.delete res = self.bigtable_hook_no_default_project_id.delete_instance( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=CBT_INSTANCE) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() delete_method.assert_called_once_with() get_client.assert_called_once_with(project_id='example-project') self.assertIsNone(res) @mock.patch( 'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch('google.cloud.bigtable.instance.Instance.create') @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_create_instance_missing_project_id(self, get_client, instance_create, mock_project_id): operation = mock.Mock() operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client) instance_create.return_value = operation with self.assertRaises(AirflowException) as cm: self.bigtable_hook_no_default_project_id.create_instance( instance_id=CBT_INSTANCE, main_cluster_id=CBT_CLUSTER, main_cluster_zone=CBT_ZONE) get_client.assert_not_called() instance_create.assert_not_called() err = cm.exception self.assertIn("The project id must be passed", str(err)) @mock.patch('google.cloud.bigtable.instance.Instance.create') @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_create_instance_overridden_project_id(self, get_client, instance_create): operation = mock.Mock() operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client) instance_create.return_value = operation res = self.bigtable_hook_no_default_project_id.create_instance( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=CBT_INSTANCE, main_cluster_id=CBT_CLUSTER, main_cluster_zone=CBT_ZONE) get_client.assert_called_once_with(project_id='example-project') instance_create.assert_called_once_with(clusters=mock.ANY) self.assertEqual(res.instance_id, 'instance') @mock.patch( 'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=None) @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_delete_table_missing_project_id(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists table_delete_method = instance_method.return_value.table.return_value.delete instance_exists_method.return_value = True with self.assertRaises(AirflowException) as cm: self.bigtable_hook_no_default_project_id.delete_table( instance_id=CBT_INSTANCE, table_id=CBT_TABLE) get_client.assert_not_called() instance_exists_method.assert_not_called() table_delete_method.assert_not_called() err = cm.exception self.assertIn("The project id must be passed", str(err)) @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_delete_table_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists table_delete_method = instance_method.return_value.table.return_value.delete instance_exists_method.return_value = True self.bigtable_hook_no_default_project_id.delete_table( project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=CBT_INSTANCE, table_id=CBT_TABLE) get_client.assert_called_once_with(project_id='example-project') instance_exists_method.assert_called_once_with() table_delete_method.assert_called_once_with()
class TestBigtableHookDefaultProjectId(unittest.TestCase): def setUp(self): with mock.patch( 'airflow.providers.google.cloud.hooks.base.CloudBaseHook.__init__', new=mock_base_gcp_hook_default_project_id): self.bigtable_hook_default_project_id = BigtableHook( gcp_conn_id='test') @mock.patch( "airflow.providers.google.cloud.hooks.bigtable.BigtableHook.client_info", new_callable=mock.PropertyMock) @mock.patch( "airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_credentials" ) @mock.patch("airflow.providers.google.cloud.hooks.bigtable.Client") def test_bigtable_client_creation(self, mock_client, mock_get_creds, mock_client_info): result = self.bigtable_hook_default_project_id._get_client( GCP_PROJECT_ID_HOOK_UNIT_TEST) mock_client.assert_called_once_with( project=GCP_PROJECT_ID_HOOK_UNIT_TEST, credentials=mock_get_creds.return_value, client_info=mock_client_info.return_value, admin=True) self.assertEqual(mock_client.return_value, result) self.assertEqual(self.bigtable_hook_default_project_id._client, result) @mock.patch( 'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST) @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_get_instance(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True res = self.bigtable_hook_default_project_id.get_instance( instance_id=CBT_INSTANCE) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() get_client.assert_called_once_with(project_id='example-project') self.assertIsNotNone(res) @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_get_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True res = self.bigtable_hook_default_project_id.get_instance( project_id='new-project', instance_id=CBT_INSTANCE) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() get_client.assert_called_once_with(project_id='new-project') self.assertIsNotNone(res) @mock.patch( 'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST) @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_get_instance_no_instance(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = False res = self.bigtable_hook_default_project_id.get_instance( instance_id=CBT_INSTANCE) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() get_client.assert_called_once_with(project_id='example-project') self.assertIsNone(res) @mock.patch( 'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST) @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_delete_instance(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True delete_method = instance_method.return_value.delete res = self.bigtable_hook_default_project_id.delete_instance( instance_id=CBT_INSTANCE) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() delete_method.assert_called_once_with() get_client.assert_called_once_with(project_id='example-project') self.assertIsNone(res) @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_delete_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True delete_method = instance_method.return_value.delete res = self.bigtable_hook_default_project_id.delete_instance( project_id='new-project', instance_id=CBT_INSTANCE) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() delete_method.assert_called_once_with() get_client.assert_called_once_with(project_id='new-project') self.assertIsNone(res) @mock.patch( 'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST) @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_delete_instance_no_instance(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = False delete_method = instance_method.return_value.delete self.bigtable_hook_default_project_id.delete_instance( instance_id=CBT_INSTANCE) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() delete_method.assert_not_called() get_client.assert_called_once_with(project_id='example-project') @mock.patch( 'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST) @mock.patch('google.cloud.bigtable.instance.Instance.create') @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_create_instance(self, get_client, instance_create, mock_project_id): operation = mock.Mock() operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client) instance_create.return_value = operation res = self.bigtable_hook_default_project_id.create_instance( instance_id=CBT_INSTANCE, main_cluster_id=CBT_CLUSTER, main_cluster_zone=CBT_ZONE) get_client.assert_called_once_with(project_id='example-project') instance_create.assert_called_once_with(clusters=mock.ANY) self.assertEqual(res.instance_id, 'instance') @mock.patch('google.cloud.bigtable.instance.Instance.create') @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_create_instance_overridden_project_id(self, get_client, instance_create): operation = mock.Mock() operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client) instance_create.return_value = operation res = self.bigtable_hook_default_project_id.create_instance( project_id='new-project', instance_id=CBT_INSTANCE, main_cluster_id=CBT_CLUSTER, main_cluster_zone=CBT_ZONE) get_client.assert_called_once_with(project_id='new-project') instance_create.assert_called_once_with(clusters=mock.ANY) self.assertEqual(res.instance_id, 'instance') @mock.patch( 'airflow.providers.google.cloud.hooks.base.CloudBaseHook.project_id', new_callable=PropertyMock, return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST) @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_delete_table(self, get_client, mock_project_id): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists table_delete_method = instance_method.return_value.table.return_value.delete instance_exists_method.return_value = True self.bigtable_hook_default_project_id.delete_table( instance_id=CBT_INSTANCE, table_id=CBT_TABLE) get_client.assert_called_once_with(project_id='example-project') instance_exists_method.assert_called_once_with() table_delete_method.assert_called_once_with() @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_delete_table_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists table_delete_method = instance_method.return_value.table.return_value.delete instance_exists_method.return_value = True self.bigtable_hook_default_project_id.delete_table( project_id='new-project', instance_id=CBT_INSTANCE, table_id=CBT_TABLE) get_client.assert_called_once_with(project_id='new-project') instance_exists_method.assert_called_once_with() table_delete_method.assert_called_once_with() @mock.patch('google.cloud.bigtable.table.Table.create') @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_create_table(self, get_client, create): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True client = mock.Mock(Client) instance = google.cloud.bigtable.instance.Instance( instance_id=CBT_INSTANCE, client=client) self.bigtable_hook_default_project_id.create_table(instance=instance, table_id=CBT_TABLE) get_client.assert_not_called() create.assert_called_once_with([], {}) @mock.patch('google.cloud.bigtable.cluster.Cluster.update') @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_update_cluster(self, get_client, update): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True client = mock.Mock(Client) instance = google.cloud.bigtable.instance.Instance( instance_id=CBT_INSTANCE, client=client) self.bigtable_hook_default_project_id.update_cluster( instance=instance, cluster_id=CBT_CLUSTER, nodes=4) get_client.assert_not_called() update.assert_called_once_with() @mock.patch('google.cloud.bigtable.table.Table.list_column_families') @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_list_column_families(self, get_client, list_column_families): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True client = mock.Mock(Client) get_client.return_value = client instance = google.cloud.bigtable.instance.Instance( instance_id=CBT_INSTANCE, client=client) self.bigtable_hook_default_project_id.get_column_families_for_table( instance=instance, table_id=CBT_TABLE) get_client.assert_not_called() list_column_families.assert_called_once_with() @mock.patch('google.cloud.bigtable.table.Table.get_cluster_states') @mock.patch( 'airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client' ) def test_get_cluster_states(self, get_client, get_cluster_states): instance_method = get_client.return_value.instance instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True client = mock.Mock(Client) instance = google.cloud.bigtable.instance.Instance( instance_id=CBT_INSTANCE, client=client) self.bigtable_hook_default_project_id.get_cluster_states_for_table( instance=instance, table_id=CBT_TABLE) get_client.assert_not_called() get_cluster_states.assert_called_once_with()