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()
