def test_do_mutate_retryable_rows(self): from google.cloud.bigtable.row import DirectRow from tests.unit._testing import _FakeStub # Setup: # - Mutate 2 rows. # Action: # - Initial attempt will mutate all 2 rows. # Expectation: # - Expect [success, non-retryable] client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_table(self.TABLE_ID, instance) row_1 = DirectRow(row_key=b'row_key', table=table) row_1.set_cell('cf', b'col', b'value1') row_2 = DirectRow(row_key=b'row_key_2', table=table) row_2.set_cell('cf', b'col', b'value2') response = self._make_responses([self.SUCCESS, self.NON_RETRYABLE]) # Patch the stub used by the API method. client._data_stub = _FakeStub([response]) worker = self._make_worker(table._instance._client, table.name, [row_1, row_2]) statuses = worker._do_mutate_retryable_rows() result = [status.code for status in statuses] expected_result = [self.SUCCESS, self.NON_RETRYABLE] self.assertEqual(result, expected_result)
def test_do_mutate_retryable_rows_mismatch_num_responses(self): from google.cloud.bigtable._generated.bigtable_pb2 import ( MutateRowsResponse) from google.cloud.bigtable.row import DirectRow from google.rpc.status_pb2 import Status from tests.unit._testing import _FakeStub client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_table(self.TABLE_ID, instance) row_1 = DirectRow(row_key=b'row_key', table=table) row_1.set_cell('cf', b'col', b'value1') row_2 = DirectRow(row_key=b'row_key_2', table=table) row_2.set_cell('cf', b'col', b'value2') response = MutateRowsResponse(entries=[ MutateRowsResponse.Entry( index=0, status=Status(code=0), ), ], ) # Patch the stub used by the API method. client._data_stub = _FakeStub([response]) worker = self._make_worker(table._instance._client, table.name, [row_1, row_2]) with self.assertRaises(RuntimeError): worker._do_mutate_retryable_rows()
def test_delete(self): from google.protobuf import empty_pb2 from tests.unit._testing import _FakeStub client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_one(self.TABLE_ID, instance) # Create request_pb request_pb = _DeleteTableRequestPB(name=self.TABLE_NAME) # Create response_pb response_pb = empty_pb2.Empty() # Patch the stub used by the API method. client._table_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # delete() has no return value. # Perform the method and check the result. result = table.delete() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'DeleteTable', (request_pb,), {}, )])
def test_do_mutate_retryable_rows(self): from google.cloud.bigtable.row import DirectRow from tests.unit._testing import _FakeStub # Setup: # - Mutate 2 rows. # Action: # - Initial attempt will mutate all 2 rows. # Expectation: # - Expect [success, non-retryable] client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_table(self.TABLE_ID, instance) row_1 = DirectRow(row_key=b'row_key', table=table) row_1.set_cell('cf', b'col', b'value1') row_2 = DirectRow(row_key=b'row_key_2', table=table) row_2.set_cell('cf', b'col', b'value2') response = self._make_responses([self.SUCCESS, self.NON_RETRYABLE]) # Patch the stub used by the API method. client._data_stub = _FakeStub([response]) worker = self._make_worker( table._instance._client, table.name, [row_1, row_2]) statuses = worker._do_mutate_retryable_rows() result = [status.code for status in statuses] expected_result = [self.SUCCESS, self.NON_RETRYABLE] self.assertEqual(result, expected_result)
def test_sample_row_keys(self): from tests.unit._testing import _FakeStub client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_one(self.TABLE_ID, instance) # Create request_pb request_pb = _SampleRowKeysRequestPB(table_name=self.TABLE_NAME) # Create response_iterator response_iterator = object() # Just passed to a mock. # Patch the stub used by the API method. client._data_stub = stub = _FakeStub(response_iterator) # Create expected_result. expected_result = response_iterator # Perform the method and check the result. result = table.sample_row_keys() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'SampleRowKeys', (request_pb,), {}, )])
def test_yield_rows(self): from tests.unit._testing import _FakeStub client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_one(self.TABLE_ID, instance) # Create response_iterator chunk = _ReadRowsResponseCellChunkPB( row_key=self.ROW_KEY, family_name=self.FAMILY_NAME, qualifier=self.QUALIFIER, timestamp_micros=self.TIMESTAMP_MICROS, value=self.VALUE, commit_row=True, ) chunks = [chunk] response = _ReadRowsResponseV2(chunks) response_iterator = _MockCancellableIterator(response) # Patch the stub used by the API method. client._data_stub = _FakeStub(response_iterator) rows = [] for row in table.yield_rows(): rows.append(row) result = rows[0] self.assertEqual(result.row_key, self.ROW_KEY)
def test_create_w_explicit_serve_nodes(self): from google.api_core import operation from google.longrunning import operations_pb2 from tests.unit._testing import _FakeStub from google.cloud.bigtable_admin_v2.gapic import ( bigtable_instance_admin_client) api = bigtable_instance_admin_client.BigtableInstanceAdminClient( mock.Mock()) credentials = _make_credentials() client = self._make_client(project=self.PROJECT, credentials=credentials, admin=True) instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID) # Create response_pb response_pb = operations_pb2.Operation(name=self.OP_NAME) # Patch the stub used by the API method. stub = _FakeStub(response_pb) client._instance_admin_client = api client._instance_admin_client.bigtable_instance_admin_stub = stub # Perform the method and check the result. result = instance.create() self.assertIsInstance(result, operation.Operation)
def test_update(self): from google.cloud.bigtable._generated import (instance_pb2 as data_v2_pb2) from tests.unit._testing import _FakeStub client = _Client(self.PROJECT) instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID, display_name=self.DISPLAY_NAME) # Create request_pb request_pb = data_v2_pb2.Instance( name=self.INSTANCE_NAME, display_name=self.DISPLAY_NAME, ) # Create response_pb response_pb = data_v2_pb2.Instance() # Patch the stub used by the API method. client._instance_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # Perform the method and check the result. result = instance.update() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'UpdateInstance', (request_pb, ), {}, )])
def test_delete(self): from google.protobuf import empty_pb2 from tests.unit._testing import _FakeStub client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_one(self.TABLE_ID, instance) # Create request_pb request_pb = _DeleteTableRequestPB(name=self.TABLE_NAME) # Create response_pb response_pb = empty_pb2.Empty() # Patch the stub used by the API method. client._table_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # delete() has no return value. # Perform the method and check the result. result = table.delete() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'DeleteTable', (request_pb, ), {}, )])
def _list_column_families_helper(self): from tests.unit._testing import _FakeStub client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_one(self.TABLE_ID, instance) # Create request_pb request_pb = _GetTableRequestPB(name=self.TABLE_NAME) # Create response_pb COLUMN_FAMILY_ID = 'foo' column_family = _ColumnFamilyPB() response_pb = _TablePB( column_families={COLUMN_FAMILY_ID: column_family}, ) # Patch the stub used by the API method. client._table_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = { COLUMN_FAMILY_ID: table.column_family(COLUMN_FAMILY_ID), } # Perform the method and check the result. result = table.list_column_families() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'GetTable', (request_pb, ), {}, )])
def test_delete(self): from google.protobuf import empty_pb2 from google.cloud.bigtable._generated import ( bigtable_instance_admin_pb2 as messages_v2_pb) from tests.unit._testing import _FakeStub client = _Client(self.PROJECT) instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID) # Create request_pb request_pb = messages_v2_pb.DeleteInstanceRequest( name=self.INSTANCE_NAME) # Create response_pb response_pb = empty_pb2.Empty() # Patch the stub used by the API method. client._instance_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # delete() has no return value. # Perform the method and check the result. result = instance.delete() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'DeleteInstance', (request_pb,), {}, )])
def test_delete(self): from google.protobuf import empty_pb2 from google.cloud.bigtable._generated import ( bigtable_instance_admin_pb2 as messages_v2_pb) from tests.unit._testing import _FakeStub client = _Client(self.PROJECT) instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID) # Create request_pb request_pb = messages_v2_pb.DeleteInstanceRequest( name=self.INSTANCE_NAME) # Create response_pb response_pb = empty_pb2.Empty() # Patch the stub used by the API method. client._instance_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # delete() has no return value. # Perform the method and check the result. result = instance.delete() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'DeleteInstance', (request_pb, ), {}, )])
def test_update(self): from google.cloud.bigtable._generated import ( instance_pb2 as data_v2_pb2) from tests.unit._testing import _FakeStub client = _Client(self.PROJECT) instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID, display_name=self.DISPLAY_NAME) # Create request_pb request_pb = data_v2_pb2.Instance( name=self.INSTANCE_NAME, display_name=self.DISPLAY_NAME, ) # Create response_pb response_pb = data_v2_pb2.Instance() # Patch the stub used by the API method. client._instance_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # Perform the method and check the result. result = instance.update() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'UpdateInstance', (request_pb,), {}, )])
def test_mutate_rows(self): from google.cloud.bigtable._generated.bigtable_pb2 import ( MutateRowsResponse) from google.cloud.bigtable.row import DirectRow from google.rpc.status_pb2 import Status from tests.unit._testing import _FakeStub client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_one(self.TABLE_ID, instance) row_1 = DirectRow(row_key=b'row_key', table=table) row_1.set_cell('cf', b'col', b'value1') row_2 = DirectRow(row_key=b'row_key_2', table=table) row_2.set_cell('cf', b'col', b'value2') response = MutateRowsResponse(entries=[ MutateRowsResponse.Entry( index=0, status=Status(code=0), ), MutateRowsResponse.Entry( index=1, status=Status(code=1), ), ], ) # Patch the stub used by the API method. client._data_stub = _FakeStub([response]) statuses = table.mutate_rows([row_1, row_2]) result = [status.code for status in statuses] expected_result = [0, 1] self.assertEqual(result, expected_result)
def test_sample_row_keys(self): from tests.unit._testing import _FakeStub client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_one(self.TABLE_ID, instance) # Create request_pb request_pb = _SampleRowKeysRequestPB(table_name=self.TABLE_NAME) # Create response_iterator response_iterator = object() # Just passed to a mock. # Patch the stub used by the API method. client._data_stub = stub = _FakeStub(response_iterator) # Create expected_result. expected_result = response_iterator # Perform the method and check the result. result = table.sample_row_keys() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'SampleRowKeys', (request_pb, ), {}, )])
def _list_column_families_helper(self): from tests.unit._testing import _FakeStub client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_one(self.TABLE_ID, instance) # Create request_pb request_pb = _GetTableRequestPB(name=self.TABLE_NAME) # Create response_pb COLUMN_FAMILY_ID = 'foo' column_family = _ColumnFamilyPB() response_pb = _TablePB( column_families={COLUMN_FAMILY_ID: column_family}, ) # Patch the stub used by the API method. client._table_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = { COLUMN_FAMILY_ID: table.column_family(COLUMN_FAMILY_ID), } # Perform the method and check the result. result = table.list_column_families() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'GetTable', (request_pb,), {}, )])
def test_do_mutate_retryable_rows_second_retry(self): from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.table import _BigtableRetryableError from tests.unit._testing import _FakeStub # Setup: # - Mutate 4 rows. # - First try results: # [success, retryable, non-retryable, retryable] # Action: # - Second try should re-attempt the 'retryable' rows. # Expectation: # - After second try: # [success, success, non-retryable, retryable] # - One of the rows tried second time returns retryable error code, # so expect a raise. # - Exception contains response whose index should be '3' even though # only two rows were retried. client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_table(self.TABLE_ID, instance) row_1 = DirectRow(row_key=b'row_key', table=table) row_1.set_cell('cf', b'col', b'value1') row_2 = DirectRow(row_key=b'row_key_2', table=table) row_2.set_cell('cf', b'col', b'value2') row_3 = DirectRow(row_key=b'row_key_3', table=table) row_3.set_cell('cf', b'col', b'value3') row_4 = DirectRow(row_key=b'row_key_4', table=table) row_4.set_cell('cf', b'col', b'value4') response = self._make_responses([self.SUCCESS, self.RETRYABLE_1]) # Patch the stub used by the API method. client._data_stub = _FakeStub([response]) worker = self._make_worker( table._instance._client, table.name, [row_1, row_2, row_3, row_4]) worker.responses_statuses = self._make_responses_statuses([ self.SUCCESS, self.RETRYABLE_1, self.NON_RETRYABLE, self.RETRYABLE_2]) with self.assertRaises(_BigtableRetryableError): worker._do_mutate_retryable_rows() statuses = worker.responses_statuses result = [status.code for status in statuses] expected_result = [self.SUCCESS, self.SUCCESS, self.NON_RETRYABLE, self.RETRYABLE_1] self.assertEqual(result, expected_result)
def test_create(self): import datetime from google.longrunning import operations_pb2 from google.protobuf.any_pb2 import Any from google.cloud.bigtable._generated import ( bigtable_instance_admin_pb2 as messages_v2_pb2) from google.cloud._helpers import _datetime_to_pb_timestamp from tests.unit._testing import _FakeStub from google.cloud.operation import Operation from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) client = _Client(self.PROJECT) instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID, display_name=self.DISPLAY_NAME) # Create response_pb metadata = messages_v2_pb2.CreateInstanceMetadata(request_time=NOW_PB) type_url = 'type.googleapis.com/%s' % ( messages_v2_pb2.CreateInstanceMetadata.DESCRIPTOR.full_name,) response_pb = operations_pb2.Operation( name=self.OP_NAME, metadata=Any( type_url=type_url, value=metadata.SerializeToString(), ) ) # Patch the stub used by the API method. client._instance_stub = stub = _FakeStub(response_pb) # Perform the method and check the result. result = instance.create() self.assertIsInstance(result, Operation) self.assertEqual(result.name, self.OP_NAME) self.assertIs(result.target, instance) self.assertIs(result.client, client) self.assertIsInstance(result.metadata, messages_v2_pb2.CreateInstanceMetadata) self.assertEqual(result.metadata.request_time, NOW_PB) self.assertEqual(result.caller_metadata, {'request_type': 'CreateInstance'}) self.assertEqual(len(stub.method_calls), 1) api_name, args, kwargs = stub.method_calls[0] self.assertEqual(api_name, 'CreateInstance') request_pb, = args self.assertIsInstance(request_pb, messages_v2_pb2.CreateInstanceRequest) self.assertEqual(request_pb.parent, 'projects/%s' % (self.PROJECT,)) self.assertEqual(request_pb.instance_id, self.INSTANCE_ID) self.assertEqual(request_pb.instance.display_name, self.DISPLAY_NAME) cluster = request_pb.clusters[self.INSTANCE_ID] self.assertEqual(cluster.serve_nodes, DEFAULT_SERVE_NODES) self.assertEqual(kwargs, {})
def test_do_mutate_retryable_rows_retry(self): from google.api_core.exceptions import ServiceUnavailable from google.cloud.bigtable._generated.bigtable_pb2 import ( MutateRowsResponse) from google.cloud.bigtable.row import DirectRow from google.rpc.status_pb2 import Status from tests.unit._testing import _FakeStub # Setup: # - Mutate 3 rows. # Action: # - Initial attempt will mutate all 3 rows. # Expectation: # - Second row returns retryable error code, so expect a raise. # - State of responses_statuses should be # [success, retryable, non-retryable] client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_table(self.TABLE_ID, instance) row_1 = DirectRow(row_key=b'row_key', table=table) row_1.set_cell('cf', b'col', b'value1') row_2 = DirectRow(row_key=b'row_key_2', table=table) row_2.set_cell('cf', b'col', b'value2') row_3 = DirectRow(row_key=b'row_key_3', table=table) row_3.set_cell('cf', b'col', b'value3') response = MutateRowsResponse(entries=[ MutateRowsResponse.Entry( index=0, status=Status(code=0), ), MutateRowsResponse.Entry( index=1, status=Status(code=4), ), MutateRowsResponse.Entry( index=2, status=Status(code=1), ), ], ) # Patch the stub used by the API method. client._data_stub = _FakeStub([response]) worker = self._make_worker(table._instance._client, table.name, [row_1, row_2, row_3]) with self.assertRaises(ServiceUnavailable): worker._do_mutate_retryable_rows() statuses = worker.responses_statuses result = [status.code for status in statuses] expected_result = [0, 4, 1] self.assertEqual(result, expected_result)
def test_list_clusters(self): from google.cloud.bigtable._generated import ( instance_pb2 as instance_v2_pb2) from google.cloud.bigtable._generated import ( bigtable_instance_admin_pb2 as messages_v2_pb2) from tests.unit._testing import _FakeStub FAILED_LOCATION = 'FAILED' FAILED_LOCATIONS = [FAILED_LOCATION] CLUSTER_ID1 = 'cluster-id1' CLUSTER_ID2 = 'cluster-id2' SERVE_NODES = 4 client = _Client(self.PROJECT) instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID) CLUSTER_NAME1 = (instance.name + '/clusters/' + CLUSTER_ID1) CLUSTER_NAME2 = (instance.name + '/clusters/' + CLUSTER_ID2) # Create request_pb request_pb = messages_v2_pb2.ListClustersRequest( parent=instance.name, ) # Create response_pb response_pb = messages_v2_pb2.ListClustersResponse( failed_locations=[FAILED_LOCATION], clusters=[ instance_v2_pb2.Cluster( name=CLUSTER_NAME1, serve_nodes=SERVE_NODES, ), instance_v2_pb2.Cluster( name=CLUSTER_NAME2, serve_nodes=SERVE_NODES, ), ], ) # Patch the stub used by the API method. client._instance_stub = stub = _FakeStub(response_pb) # Create expected_result. clusters = [ instance.cluster(CLUSTER_ID1), instance.cluster(CLUSTER_ID2), ] expected_result = (clusters, FAILED_LOCATIONS) # Perform the method and check the result. result = instance.list_clusters() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'ListClusters', (request_pb,), {}, )])
def test_commit(self): from google.cloud._testing import _Monkey from tests.unit._testing import _FakeStub from google.cloud.bigtable import row as MUT row_key = b'row_key' table_name = 'projects/more-stuff' column_family_id = u'column_family_id' column = b'column' client = _Client() table = _Table(table_name, client=client) row = self._make_one(row_key, table) # Create request_pb value = b'bytes-value' # We will call row.append_cell_value(COLUMN_FAMILY_ID, COLUMN, value). request_pb = _ReadModifyWriteRowRequestPB( table_name=table_name, row_key=row_key, rules=[ _ReadModifyWriteRulePB( family_name=column_family_id, column_qualifier=column, append_value=value, ), ], ) # Create response_pb response_pb = object() # Patch the stub used by the API method. client._data_stub = stub = _FakeStub(response_pb) # Create expected_result. row_responses = [] expected_result = object() def mock_parse_rmw_row_response(row_response): row_responses.append(row_response) return expected_result # Perform the method and check the result. with _Monkey(MUT, _parse_rmw_row_response=mock_parse_rmw_row_response): row.append_cell_value(column_family_id, column, value) result = row.commit() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'ReadModifyWriteRow', (request_pb, ), {}, )]) self.assertEqual(row_responses, [response_pb]) self.assertEqual(row._rule_pb_list, [])
def test_do_mutate_retryable_rows_second_try(self): from google.cloud.bigtable._generated.bigtable_pb2 import ( MutateRowsResponse) from google.cloud.bigtable.row import DirectRow from google.rpc.status_pb2 import Status from tests.unit._testing import _FakeStub # Setup: # - Mutate 4 rows. # - First try results: # [success, retryable, non-retryable, retryable] # Action: # - Second try should re-attempt the 'retryable' rows. # Expectation: # - After second try: # [success, non-retryable, non-retryable, success] client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_table(self.TABLE_ID, instance) row_1 = DirectRow(row_key=b'row_key', table=table) row_1.set_cell('cf', b'col', b'value1') row_2 = DirectRow(row_key=b'row_key_2', table=table) row_2.set_cell('cf', b'col', b'value2') row_3 = DirectRow(row_key=b'row_key_3', table=table) row_3.set_cell('cf', b'col', b'value3') row_4 = DirectRow(row_key=b'row_key_4', table=table) row_4.set_cell('cf', b'col', b'value4') response = MutateRowsResponse(entries=[ MutateRowsResponse.Entry( index=0, status=Status(code=1), ), MutateRowsResponse.Entry( index=1, status=Status(code=0), ), ], ) # Patch the stub used by the API method. client._data_stub = _FakeStub([response]) worker = self._make_worker(table._instance._client, table.name, [row_1, row_2, row_3, row_4]) worker.responses_statuses = self._make_responses_statuses( [0, 4, 1, 10]) statuses = worker._do_mutate_retryable_rows() result = [status.code for status in statuses] expected_result = [0, 1, 1, 0] self.assertEqual(result, expected_result)
def test_commit(self): from google.cloud._testing import _Monkey from tests.unit._testing import _FakeStub from google.cloud.bigtable import row as MUT row_key = b'row_key' table_name = 'projects/more-stuff' column_family_id = u'column_family_id' column = b'column' client = _Client() table = _Table(table_name, client=client) row = self._make_one(row_key, table) # Create request_pb value = b'bytes-value' # We will call row.append_cell_value(COLUMN_FAMILY_ID, COLUMN, value). request_pb = _ReadModifyWriteRowRequestPB( table_name=table_name, row_key=row_key, rules=[ _ReadModifyWriteRulePB( family_name=column_family_id, column_qualifier=column, append_value=value, ), ], ) # Create response_pb response_pb = object() # Patch the stub used by the API method. client._data_stub = stub = _FakeStub(response_pb) # Create expected_result. row_responses = [] expected_result = object() def mock_parse_rmw_row_response(row_response): row_responses.append(row_response) return expected_result # Perform the method and check the result. with _Monkey(MUT, _parse_rmw_row_response=mock_parse_rmw_row_response): row.append_cell_value(column_family_id, column, value) result = row.commit() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'ReadModifyWriteRow', (request_pb,), {}, )]) self.assertEqual(row_responses, [response_pb]) self.assertEqual(row._rule_pb_list, [])
def _create_test_helper(gc_rule=None): from google.cloud.bigtable_admin_v2.types import ( bigtable_table_admin as table_admin_v2_pb2, ) from tests.unit._testing import _FakeStub from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import ( BigtableTableAdminClient, ) project_id = "project-id" zone = "zone" cluster_id = "cluster-id" table_id = "table-id" column_family_id = "column-family-id" table_name = ("projects/" + project_id + "/zones/" + zone + "/clusters/" + cluster_id + "/tables/" + table_id) api = mock.create_autospec(BigtableTableAdminClient) credentials = _make_credentials() client = _make_client(project=project_id, credentials=credentials, admin=True) table = _Table(table_name, client=client) column_family = _make_column_family(column_family_id, table, gc_rule=gc_rule) # Create request_pb if gc_rule is None: column_family_pb = _ColumnFamilyPB() else: column_family_pb = _ColumnFamilyPB(gc_rule=gc_rule.to_pb()) request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest( name=table_name) modification = table_admin_v2_pb2.ModifyColumnFamiliesRequest.Modification( ) modification.id = column_family_id modification.create = column_family_pb request_pb.modifications.append(modification) # Create response_pb response_pb = _ColumnFamilyPB() # Patch the stub used by the API method. stub = _FakeStub(response_pb) client._table_admin_client = api client._table_admin_client.transport.create = stub # Create expected_result. expected_result = None # create() has no return value. # Perform the method and check the result. assert stub.results == (response_pb, ) result = column_family.create() assert result == expected_result
def test_do_mutate_retryable_rows_second_retry(self): from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.table import _BigtableRetryableError from tests.unit._testing import _FakeStub # Setup: # - Mutate 4 rows. # - First try results: # [success, retryable, non-retryable, retryable] # Action: # - Second try should re-attempt the 'retryable' rows. # Expectation: # - After second try: # [success, success, non-retryable, retryable] # - One of the rows tried second time returns retryable error code, # so expect a raise. # - Exception contains response whose index should be '3' even though # only two rows were retried. client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_table(self.TABLE_ID, instance) row_1 = DirectRow(row_key=b'row_key', table=table) row_1.set_cell('cf', b'col', b'value1') row_2 = DirectRow(row_key=b'row_key_2', table=table) row_2.set_cell('cf', b'col', b'value2') row_3 = DirectRow(row_key=b'row_key_3', table=table) row_3.set_cell('cf', b'col', b'value3') row_4 = DirectRow(row_key=b'row_key_4', table=table) row_4.set_cell('cf', b'col', b'value4') response = self._make_responses([self.SUCCESS, self.RETRYABLE_1]) # Patch the stub used by the API method. client._data_stub = _FakeStub([response]) worker = self._make_worker(table._instance._client, table.name, [row_1, row_2, row_3, row_4]) worker.responses_statuses = self._make_responses_statuses([ self.SUCCESS, self.RETRYABLE_1, self.NON_RETRYABLE, self.RETRYABLE_2 ]) with self.assertRaises(_BigtableRetryableError): worker._do_mutate_retryable_rows() statuses = worker.responses_statuses result = [status.code for status in statuses] expected_result = [ self.SUCCESS, self.SUCCESS, self.NON_RETRYABLE, self.RETRYABLE_1 ] self.assertEqual(result, expected_result)
def _update_test_helper(self, gc_rule=None): from tests.unit._testing import _FakeStub from google.cloud.bigtable_admin_v2.proto import ( bigtable_table_admin_pb2 as table_admin_v2_pb2, ) from google.cloud.bigtable_admin_v2.gapic import bigtable_table_admin_client project_id = "project-id" zone = "zone" cluster_id = "cluster-id" table_id = "table-id" column_family_id = "column-family-id" table_name = ( "projects/" + project_id + "/zones/" + zone + "/clusters/" + cluster_id + "/tables/" + table_id ) api = bigtable_table_admin_client.BigtableTableAdminClient(mock.Mock()) credentials = _make_credentials() client = self._make_client( project=project_id, credentials=credentials, admin=True ) table = _Table(table_name, client=client) column_family = self._make_one(column_family_id, table, gc_rule=gc_rule) # Create request_pb if gc_rule is None: column_family_pb = _ColumnFamilyPB() else: column_family_pb = _ColumnFamilyPB(gc_rule=gc_rule.to_pb()) request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(name=table_name) request_pb.modifications.add(id=column_family_id, update=column_family_pb) # Create response_pb response_pb = _ColumnFamilyPB() # Patch the stub used by the API method. stub = _FakeStub(response_pb) client._table_admin_client = api client._table_admin_client.transport.update = stub # Create expected_result. expected_result = None # update() has no return value. # Perform the method and check the result. self.assertEqual(stub.results, (response_pb,)) result = column_family.update() self.assertEqual(result, expected_result)
def test_create(self): import datetime from google.api_core import operation from google.longrunning import operations_pb2 from google.protobuf.any_pb2 import Any from google.cloud.bigtable._generated import ( bigtable_instance_admin_pb2 as messages_v2_pb2) from google.cloud._helpers import _datetime_to_pb_timestamp from tests.unit._testing import _FakeStub from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) client = _Client(self.PROJECT) instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID, display_name=self.DISPLAY_NAME) # Create response_pb metadata = messages_v2_pb2.CreateInstanceMetadata(request_time=NOW_PB) type_url = 'type.googleapis.com/%s' % ( messages_v2_pb2.CreateInstanceMetadata.DESCRIPTOR.full_name,) response_pb = operations_pb2.Operation( name=self.OP_NAME, metadata=Any( type_url=type_url, value=metadata.SerializeToString(), ) ) # Patch the stub used by the API method. client._instance_stub = stub = _FakeStub(response_pb) # Perform the method and check the result. result = instance.create() self.assertIsInstance(result, operation.Operation) self.assertEqual(result.operation.name, self.OP_NAME) self.assertIsInstance(result.metadata, messages_v2_pb2.CreateInstanceMetadata) self.assertEqual(result.metadata.request_time, NOW_PB) self.assertEqual(len(stub.method_calls), 1) api_name, args, kwargs = stub.method_calls[0] self.assertEqual(api_name, 'CreateInstance') request_pb, = args self.assertIsInstance(request_pb, messages_v2_pb2.CreateInstanceRequest) self.assertEqual(request_pb.parent, 'projects/%s' % (self.PROJECT,)) self.assertEqual(request_pb.instance_id, self.INSTANCE_ID) self.assertEqual(request_pb.instance.display_name, self.DISPLAY_NAME) cluster = request_pb.clusters[self.INSTANCE_ID] self.assertEqual(cluster.serve_nodes, DEFAULT_SERVE_NODES) self.assertEqual(kwargs, {})
def test_read_rows(self): from google.cloud._testing import _Monkey from tests.unit._testing import _FakeStub from google.cloud.bigtable.row_data import PartialRowsData from google.cloud.bigtable import table as MUT client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_one(self.TABLE_ID, instance) # Create request_pb request_pb = object() # Returned by our mock. mock_created = [] def mock_create_row_request(table_name, **kwargs): mock_created.append((table_name, kwargs)) return request_pb # Create response_iterator response_iterator = object() # Patch the stub used by the API method. client._data_stub = stub = _FakeStub(response_iterator) # Create expected_result. expected_result = PartialRowsData(response_iterator) # Perform the method and check the result. start_key = b'start-key' end_key = b'end-key' filter_obj = object() limit = 22 with _Monkey(MUT, _create_row_request=mock_create_row_request): result = table.read_rows(start_key=start_key, end_key=end_key, filter_=filter_obj, limit=limit) self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'ReadRows', (request_pb, ), {}, )]) created_kwargs = { 'start_key': start_key, 'end_key': end_key, 'filter_': filter_obj, 'limit': limit, 'end_inclusive': False, } self.assertEqual(mock_created, [(table.name, created_kwargs)])
def test_update_app_profile_with_single_routing_policy(self): import datetime from google.api_core import operation from google.longrunning import operations_pb2 from google.protobuf.any_pb2 import Any from google.cloud.bigtable_admin_v2.proto import ( bigtable_instance_admin_pb2 as messages_v2_pb2) from google.cloud._helpers import _datetime_to_pb_timestamp from tests.unit._testing import _FakeStub from google.cloud.bigtable_admin_v2.gapic import ( bigtable_instance_admin_client) instance_api = ( bigtable_instance_admin_client.BigtableInstanceAdminClient( mock.Mock())) credentials = _make_credentials() client = self._make_client(project=self.PROJECT, credentials=credentials, admin=True) instance = self._make_one(self.INSTANCE_ID, client) # Create response_pb NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) metadata = messages_v2_pb2.CreateInstanceMetadata(request_time=NOW_PB) type_url = 'type.googleapis.com/%s' % ( messages_v2_pb2.CreateInstanceMetadata.DESCRIPTOR.full_name, ) response_pb = operations_pb2.Operation( name=self.OP_NAME, metadata=Any( type_url=type_url, value=metadata.SerializeToString(), )) # Patch the stub used by the API method. client._instance_admin_client = instance_api stub = _FakeStub(response_pb) client._instance_admin_client.bigtable_instance_admin_stub = stub update_mask = [] # Perform the method and check the result. app_profile_id = 'appProfileId1262094415' cluster_id = 'cluster-id' result = instance.update_app_profile(app_profile_id, update_mask=update_mask, routing_policy_type=2, cluster_id=cluster_id) self.assertIsInstance(result, operation.Operation)
def test_do_mutate_retryable_rows_second_try(self): from google.cloud.bigtable.row import DirectRow from tests.unit._testing import _FakeStub # Setup: # - Mutate 4 rows. # - First try results: # [success, retryable, non-retryable, retryable] # Action: # - Second try should re-attempt the 'retryable' rows. # Expectation: # - After second try: # [success, non-retryable, non-retryable, success] client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_table(self.TABLE_ID, instance) row_1 = DirectRow(row_key=b'row_key', table=table) row_1.set_cell('cf', b'col', b'value1') row_2 = DirectRow(row_key=b'row_key_2', table=table) row_2.set_cell('cf', b'col', b'value2') row_3 = DirectRow(row_key=b'row_key_3', table=table) row_3.set_cell('cf', b'col', b'value3') row_4 = DirectRow(row_key=b'row_key_4', table=table) row_4.set_cell('cf', b'col', b'value4') response = self._make_responses([self.NON_RETRYABLE, self.SUCCESS]) # Patch the stub used by the API method. client._data_stub = _FakeStub([response]) worker = self._make_worker( table._instance._client, table.name, [row_1, row_2, row_3, row_4]) worker.responses_statuses = self._make_responses_statuses([ self.SUCCESS, self.RETRYABLE_1, self.NON_RETRYABLE, self.RETRYABLE_2]) statuses = worker._do_mutate_retryable_rows() result = [status.code for status in statuses] expected_result = [self.SUCCESS, self.NON_RETRYABLE, self.NON_RETRYABLE, self.SUCCESS] self.assertEqual(result, expected_result)
def _update_test_helper(self, gc_rule=None): from tests.unit._testing import _FakeStub from google.cloud.bigtable._generated import (bigtable_table_admin_pb2 as table_admin_v2_pb2) project_id = 'project-id' zone = 'zone' cluster_id = 'cluster-id' table_id = 'table-id' column_family_id = 'column-family-id' table_name = ('projects/' + project_id + '/zones/' + zone + '/clusters/' + cluster_id + '/tables/' + table_id) client = _Client() table = _Table(table_name, client=client) column_family = self._make_one(column_family_id, table, gc_rule=gc_rule) # Create request_pb if gc_rule is None: column_family_pb = _ColumnFamilyPB() else: column_family_pb = _ColumnFamilyPB(gc_rule=gc_rule.to_pb()) request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest( name=table_name) request_pb.modifications.add( id=column_family_id, update=column_family_pb, ) # Create response_pb response_pb = _ColumnFamilyPB() # Patch the stub used by the API method. client._table_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # update() has no return value. # Perform the method and check the result. self.assertEqual(stub.results, (response_pb, )) result = column_family.update() self.assertEqual(stub.results, ()) self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'ModifyColumnFamilies', (request_pb, ), {}, )])
def _update_test_helper(self, gc_rule=None): from tests.unit._testing import _FakeStub from google.cloud.bigtable._generated import ( bigtable_table_admin_pb2 as table_admin_v2_pb2) project_id = 'project-id' zone = 'zone' cluster_id = 'cluster-id' table_id = 'table-id' column_family_id = 'column-family-id' table_name = ('projects/' + project_id + '/zones/' + zone + '/clusters/' + cluster_id + '/tables/' + table_id) client = _Client() table = _Table(table_name, client=client) column_family = self._make_one( column_family_id, table, gc_rule=gc_rule) # Create request_pb if gc_rule is None: column_family_pb = _ColumnFamilyPB() else: column_family_pb = _ColumnFamilyPB(gc_rule=gc_rule.to_pb()) request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest( name=table_name) request_pb.modifications.add( id=column_family_id, update=column_family_pb, ) # Create response_pb response_pb = _ColumnFamilyPB() # Patch the stub used by the API method. client._table_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # update() has no return value. # Perform the method and check the result. self.assertEqual(stub.results, (response_pb,)) result = column_family.update() self.assertEqual(stub.results, ()) self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'ModifyColumnFamilies', (request_pb,), {}, )])
def test_create(self): import datetime from google.api_core import operation from google.longrunning import operations_pb2 from google.protobuf.any_pb2 import Any from google.cloud.bigtable_admin_v2.proto import ( bigtable_instance_admin_pb2 as messages_v2_pb2) from google.cloud._helpers import _datetime_to_pb_timestamp from tests.unit._testing import _FakeStub from google.cloud.bigtable_admin_v2.gapic import ( bigtable_instance_admin_client) NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) api = bigtable_instance_admin_client.BigtableInstanceAdminClient( mock.Mock()) credentials = _make_credentials() client = self._make_client(project=self.PROJECT, credentials=credentials, admin=True) instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID, display_name=self.DISPLAY_NAME) # Create response_pb metadata = messages_v2_pb2.CreateInstanceMetadata(request_time=NOW_PB) type_url = 'type.googleapis.com/%s' % ( messages_v2_pb2.CreateInstanceMetadata.DESCRIPTOR.full_name, ) response_pb = operations_pb2.Operation( name=self.OP_NAME, metadata=Any( type_url=type_url, value=metadata.SerializeToString(), )) # Patch the stub used by the API method. stub = _FakeStub(response_pb) client._instance_admin_client = api client._instance_admin_client.bigtable_instance_admin_stub = stub # Perform the method and check the result. result = instance.create() self.assertIsInstance(result, operation.Operation) # self.assertEqual(result.operation.name, self.OP_NAME) self.assertIsInstance(result.metadata, messages_v2_pb2.CreateInstanceMetadata)
def test_commit(self): from google.protobuf import empty_pb2 from tests.unit._testing import _FakeStub row_key = b'row_key' table_name = 'projects/more-stuff' column_family_id = u'column_family_id' column = b'column' client = _Client() table = _Table(table_name, client=client) row = self._make_one(row_key, table) # Create request_pb value = b'bytes-value' mutation = _MutationPB( set_cell=_MutationSetCellPB( family_name=column_family_id, column_qualifier=column, timestamp_micros=-1, # Default value. value=value, ), ) request_pb = _MutateRowRequestPB( table_name=table_name, row_key=row_key, mutations=[mutation], ) # Create response_pb response_pb = empty_pb2.Empty() # Patch the stub used by the API method. client._data_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # commit() has no return value when no filter. # Perform the method and check the result. row.set_cell(column_family_id, column, value) result = row.commit() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'MutateRow', (request_pb,), {}, )]) self.assertEqual(row._pb_mutations, [])
def test_do_mutate_retryable_rows_second_try(self): from google.cloud.bigtable.row import DirectRow from tests.unit._testing import _FakeStub # Setup: # - Mutate 4 rows. # - First try results: # [success, retryable, non-retryable, retryable] # Action: # - Second try should re-attempt the 'retryable' rows. # Expectation: # - After second try: # [success, non-retryable, non-retryable, success] client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_table(self.TABLE_ID, instance) row_1 = DirectRow(row_key=b'row_key', table=table) row_1.set_cell('cf', b'col', b'value1') row_2 = DirectRow(row_key=b'row_key_2', table=table) row_2.set_cell('cf', b'col', b'value2') row_3 = DirectRow(row_key=b'row_key_3', table=table) row_3.set_cell('cf', b'col', b'value3') row_4 = DirectRow(row_key=b'row_key_4', table=table) row_4.set_cell('cf', b'col', b'value4') response = self._make_responses([self.NON_RETRYABLE, self.SUCCESS]) # Patch the stub used by the API method. client._data_stub = _FakeStub([response]) worker = self._make_worker(table._instance._client, table.name, [row_1, row_2, row_3, row_4]) worker.responses_statuses = self._make_responses_statuses([ self.SUCCESS, self.RETRYABLE_1, self.NON_RETRYABLE, self.RETRYABLE_2 ]) statuses = worker._do_mutate_retryable_rows() result = [status.code for status in statuses] expected_result = [ self.SUCCESS, self.NON_RETRYABLE, self.NON_RETRYABLE, self.SUCCESS ] self.assertEqual(result, expected_result)
def test_read_rows_one_chunk(self): from google.cloud._testing import _Monkey from tests.unit._testing import _FakeStub from google.cloud.bigtable import retry as MUT from google.cloud.bigtable.retry import ReadRowsIterator from google.cloud.bigtable.row_data import Cell from google.cloud.bigtable.row_data import PartialRowsData client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_one(self.TABLE_ID, instance) # Create request_pb request_pb = object() # Returned by our mock. mock_created = [] def mock_create_row_request(table_name, **kwargs): mock_created.append((table_name, kwargs)) return request_pb # Create response_iterator chunk = _ReadRowsResponseCellChunkPB( row_key=self.ROW_KEY, family_name=self.FAMILY_NAME, qualifier=self.QUALIFIER, timestamp_micros=self.TIMESTAMP_MICROS, value=self.VALUE, commit_row=True, ) response_pb = _ReadRowsResponsePB(chunks=[chunk]) response_iterator = iter([response_pb]) # Patch the stub used by the API method. client._data_stub = stub = _FakeStub(response_iterator) start_key = b'start-key' end_key = b'end-key' filter_obj = object() limit = 22 with _Monkey(MUT, _create_row_request=mock_create_row_request): # Perform the method and check the result. result = table.read_rows(start_key=start_key, end_key=end_key, filter_=filter_obj, limit=limit) result.consume_all()
def test_commit_no_rules(self): from tests.unit._testing import _FakeStub row_key = b'row_key' client = _Client() table = _Table(None, client=client) row = self._make_one(row_key, table) self.assertEqual(row._rule_pb_list, []) # Patch the stub used by the API method. client._data_stub = stub = _FakeStub() # Perform the method and check the result. result = row.commit() self.assertEqual(result, {}) # Make sure no request was sent. self.assertEqual(stub.method_calls, [])
def test_delete(self): from google.protobuf import empty_pb2 from google.cloud.bigtable_admin_v2.types import ( bigtable_table_admin as table_admin_v2_pb2, ) from tests.unit._testing import _FakeStub from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import ( BigtableTableAdminClient, ) project_id = "project-id" zone = "zone" cluster_id = "cluster-id" table_id = "table-id" column_family_id = "column-family-id" table_name = ("projects/" + project_id + "/zones/" + zone + "/clusters/" + cluster_id + "/tables/" + table_id) api = mock.create_autospec(BigtableTableAdminClient) credentials = _make_credentials() client = self._make_client(project=project_id, credentials=credentials, admin=True) table = _Table(table_name, client=client) column_family = self._make_one(column_family_id, table) # Create request_pb request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest( name=table_name) modification = table_admin_v2_pb2.ModifyColumnFamiliesRequest.Modification( id=column_family_id, drop=True) request_pb.modifications.append(modification) # Create response_pb response_pb = empty_pb2.Empty() # Patch the stub used by the API method. stub = _FakeStub(response_pb) client._table_admin_client = api client._table_admin_client.transport.delete = stub # Create expected_result. expected_result = None # delete() has no return value. # Perform the method and check the result. self.assertEqual(stub.results, (response_pb, )) result = column_family.delete() self.assertEqual(result, expected_result)
def test_delete(self): from google.protobuf import empty_pb2 from google.cloud.bigtable_admin_v2.proto import ( bigtable_table_admin_pb2 as table_admin_v2_pb2) from tests.unit._testing import _FakeStub from google.cloud.bigtable_admin_v2.gapic import ( bigtable_table_admin_client) project_id = 'project-id' zone = 'zone' cluster_id = 'cluster-id' table_id = 'table-id' column_family_id = 'column-family-id' table_name = ('projects/' + project_id + '/zones/' + zone + '/clusters/' + cluster_id + '/tables/' + table_id) api = bigtable_table_admin_client.BigtableTableAdminClient(mock.Mock()) credentials = _make_credentials() client = self._make_client(project=project_id, credentials=credentials, admin=True) table = _Table(table_name, client=client) column_family = self._make_one(column_family_id, table) # Create request_pb request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest( name=table_name) request_pb.modifications.add( id=column_family_id, drop=True) # Create response_pb response_pb = empty_pb2.Empty() # Patch the stub used by the API method. stub = _FakeStub(response_pb) client._table_admin_client = api client._table_admin_client.bigtable_table_admin_stub = stub # Create expected_result. expected_result = None # delete() has no return value. # Perform the method and check the result. self.assertEqual(stub.results, (response_pb,)) result = column_family.delete() self.assertEqual(result, expected_result)
def test_do_mutate_retryable_rows(self): from google.cloud.bigtable._generated.bigtable_pb2 import ( MutateRowsResponse) from google.cloud.bigtable.row import DirectRow from google.rpc.status_pb2 import Status from tests.unit._testing import _FakeStub # Setup: # - Mutate 2 rows. # Action: # - Initial attempt will mutate all 2 rows. # Expectation: # - Expect [success, non-retryable] client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_table(self.TABLE_ID, instance) row_1 = DirectRow(row_key=b'row_key', table=table) row_1.set_cell('cf', b'col', b'value1') row_2 = DirectRow(row_key=b'row_key_2', table=table) row_2.set_cell('cf', b'col', b'value2') response = MutateRowsResponse(entries=[ MutateRowsResponse.Entry( index=0, status=Status(code=0), ), MutateRowsResponse.Entry( index=1, status=Status(code=1), ), ], ) # Patch the stub used by the API method. client._data_stub = _FakeStub([response]) worker = self._make_worker(table._instance._client, table.name, [row_1, row_2]) statuses = worker._do_mutate_retryable_rows() result = [status.code for status in statuses] expected_result = [0, 1] self.assertEqual(result, expected_result)
def test_delete(self): from google.protobuf import empty_pb2 from google.cloud.bigtable_admin_v2.proto import ( bigtable_table_admin_pb2 as table_admin_v2_pb2) from tests.unit._testing import _FakeStub from google.cloud.bigtable_admin_v2.gapic import ( bigtable_table_admin_client) project_id = 'project-id' zone = 'zone' cluster_id = 'cluster-id' table_id = 'table-id' column_family_id = 'column-family-id' table_name = ('projects/' + project_id + '/zones/' + zone + '/clusters/' + cluster_id + '/tables/' + table_id) api = bigtable_table_admin_client.BigtableTableAdminClient(mock.Mock()) credentials = _make_credentials() client = self._make_client(project=project_id, credentials=credentials, admin=True) table = _Table(table_name, client=client) column_family = self._make_one(column_family_id, table) # Create request_pb request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest( name=table_name) request_pb.modifications.add( id=column_family_id, drop=True) # Create response_pb response_pb = empty_pb2.Empty() # Patch the stub used by the API method. stub = _FakeStub(response_pb) client._table_admin_client = api client._table_admin_client.transport.delete = stub # Create expected_result. expected_result = None # delete() has no return value. # Perform the method and check the result. self.assertEqual(stub.results, (response_pb,)) result = column_family.delete() self.assertEqual(result, expected_result)
def _create_test_helper(self, initial_split_keys, column_families=()): from google.cloud._helpers import _to_bytes from tests.unit._testing import _FakeStub client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_one(self.TABLE_ID, instance) # Create request_pb splits_pb = [ _CreateTableRequestSplitPB(key=_to_bytes(key)) for key in initial_split_keys or () ] table_pb = None if column_families: table_pb = _TablePB() for cf in column_families: cf_pb = table_pb.column_families[cf.column_family_id] if cf.gc_rule is not None: cf_pb.gc_rule.MergeFrom(cf.gc_rule.to_pb()) request_pb = _CreateTableRequestPB( initial_splits=splits_pb, parent=self.INSTANCE_NAME, table_id=self.TABLE_ID, table=table_pb, ) # Create response_pb response_pb = _TablePB() # Patch the stub used by the API method. client._table_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # create() has no return value. # Perform the method and check the result. result = table.create(initial_split_keys=initial_split_keys, column_families=column_families) self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'CreateTable', (request_pb, ), {}, )])
def _create_test_helper(self, initial_split_keys, column_families=()): from google.cloud._helpers import _to_bytes from tests.unit._testing import _FakeStub client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_one(self.TABLE_ID, instance) # Create request_pb splits_pb = [ _CreateTableRequestSplitPB(key=_to_bytes(key)) for key in initial_split_keys or ()] table_pb = None if column_families: table_pb = _TablePB() for cf in column_families: cf_pb = table_pb.column_families[cf.column_family_id] if cf.gc_rule is not None: cf_pb.gc_rule.CopyFrom(cf.gc_rule.to_pb()) request_pb = _CreateTableRequestPB( initial_splits=splits_pb, parent=self.INSTANCE_NAME, table_id=self.TABLE_ID, table=table_pb, ) # Create response_pb response_pb = _TablePB() # Patch the stub used by the API method. client._table_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # create() has no return value. # Perform the method and check the result. result = table.create(initial_split_keys=initial_split_keys, column_families=column_families) self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'CreateTable', (request_pb,), {}, )])
def test_do_mutate_retryable_rows_retry(self): from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.table import _BigtableRetryableError from tests.unit._testing import _FakeStub # Setup: # - Mutate 3 rows. # Action: # - Initial attempt will mutate all 3 rows. # Expectation: # - Second row returns retryable error code, so expect a raise. # - State of responses_statuses should be # [success, retryable, non-retryable] client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_table(self.TABLE_ID, instance) row_1 = DirectRow(row_key=b'row_key', table=table) row_1.set_cell('cf', b'col', b'value1') row_2 = DirectRow(row_key=b'row_key_2', table=table) row_2.set_cell('cf', b'col', b'value2') row_3 = DirectRow(row_key=b'row_key_3', table=table) row_3.set_cell('cf', b'col', b'value3') response = self._make_responses([ self.SUCCESS, self.RETRYABLE_1, self.NON_RETRYABLE]) # Patch the stub used by the API method. client._data_stub = _FakeStub([response]) worker = self._make_worker( table._instance._client, table.name, [row_1, row_2, row_3]) with self.assertRaises(_BigtableRetryableError): worker._do_mutate_retryable_rows() statuses = worker.responses_statuses result = [status.code for status in statuses] expected_result = [self.SUCCESS, self.RETRYABLE_1, self.NON_RETRYABLE] self.assertEqual(result, expected_result)
def test_delete(self): from google.protobuf import empty_pb2 from google.cloud.bigtable._generated import ( bigtable_table_admin_pb2 as table_admin_v2_pb2) from tests.unit._testing import _FakeStub project_id = 'project-id' zone = 'zone' cluster_id = 'cluster-id' table_id = 'table-id' column_family_id = 'column-family-id' table_name = ('projects/' + project_id + '/zones/' + zone + '/clusters/' + cluster_id + '/tables/' + table_id) client = _Client() table = _Table(table_name, client=client) column_family = self._make_one(column_family_id, table) # Create request_pb request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest( name=table_name) request_pb.modifications.add( id=column_family_id, drop=True) # Create response_pb response_pb = empty_pb2.Empty() # Patch the stub used by the API method. client._table_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # delete() has no return value. # Perform the method and check the result. self.assertEqual(stub.results, (response_pb,)) result = column_family.delete() self.assertEqual(stub.results, ()) self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'ModifyColumnFamilies', (request_pb,), {}, )])
def test_commit_no_mutations(self): from tests.unit._testing import _FakeStub row_key = b'row_key' client = _Client() table = _Table(None, client=client) filter_ = object() row = self._make_one(row_key, table, filter_=filter_) self.assertEqual(row._true_pb_mutations, []) self.assertEqual(row._false_pb_mutations, []) # Patch the stub used by the API method. client._data_stub = stub = _FakeStub() # Perform the method and check the result. result = row.commit() self.assertIsNone(result) # Make sure no request was sent. self.assertEqual(stub.method_calls, [])
def test_commit_no_rules(self): from tests.unit._testing import _FakeStub project_id = 'project-id' row_key = b'row_key' credentials = _make_credentials() client = self._make_client(project=project_id, credentials=credentials, admin=True) table = _Table(None, client=client) row = self._make_one(row_key, table) self.assertEqual(row._rule_pb_list, []) # Patch the stub used by the API method. stub = _FakeStub() # Perform the method and check the result. result = row.commit() self.assertEqual(result, {}) # Make sure no request was sent. self.assertEqual(stub.method_calls, [])
def test_reload(self): from google.cloud.bigtable._generated import ( instance_pb2 as data_v2_pb2) from google.cloud.bigtable._generated import ( bigtable_instance_admin_pb2 as messages_v2_pb) from tests.unit._testing import _FakeStub client = _Client(self.PROJECT) instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID) # Create request_pb request_pb = messages_v2_pb.GetInstanceRequest( name=self.INSTANCE_NAME) # Create response_pb DISPLAY_NAME = u'hey-hi-hello' response_pb = data_v2_pb2.Instance( display_name=DISPLAY_NAME, ) # Patch the stub used by the API method. client._instance_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # reload() has no return value. # Check Instance optional config values before. self.assertEqual(instance.display_name, self.INSTANCE_ID) # Perform the method and check the result. result = instance.reload() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'GetInstance', (request_pb,), {}, )]) # Check Instance optional config values before. self.assertEqual(instance.display_name, DISPLAY_NAME)
def _list_tables_helper(self, table_name=None): from google.cloud.bigtable._generated import ( table_pb2 as table_data_v2_pb2) from google.cloud.bigtable._generated import ( bigtable_table_admin_pb2 as table_messages_v1_pb2) from tests.unit._testing import _FakeStub client = _Client(self.PROJECT) instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID) # Create request_ request_pb = table_messages_v1_pb2.ListTablesRequest( parent=self.INSTANCE_NAME) # Create response_pb if table_name is None: table_name = self.TABLE_NAME response_pb = table_messages_v1_pb2.ListTablesResponse( tables=[ table_data_v2_pb2.Table(name=table_name), ], ) # Patch the stub used by the API method. client._table_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_table = instance.table(self.TABLE_ID) expected_result = [expected_table] # Perform the method and check the result. result = instance.list_tables() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'ListTables', (request_pb,), {}, )])
def _read_row_helper(self, chunks, expected_result): from google.cloud._testing import _Monkey from tests.unit._testing import _FakeStub from google.cloud.bigtable import table as MUT client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_one(self.TABLE_ID, instance) # Create request_pb request_pb = object() # Returned by our mock. mock_created = [] def mock_create_row_request(table_name, row_key, filter_): mock_created.append((table_name, row_key, filter_)) return request_pb # Create response_iterator if chunks is None: response_iterator = iter(()) # no responses at all else: response_pb = _ReadRowsResponsePB(chunks=chunks) response_iterator = iter([response_pb]) # Patch the stub used by the API method. client._data_stub = stub = _FakeStub(response_iterator) # Perform the method and check the result. filter_obj = object() with _Monkey(MUT, _create_row_request=mock_create_row_request): result = table.read_row(self.ROW_KEY, filter_=filter_obj) self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'ReadRows', (request_pb,), {}, )]) self.assertEqual(mock_created, [(table.name, self.ROW_KEY, filter_obj)])
def test_create(self): from google.api_core import operation from google.longrunning import operations_pb2 from google.cloud.bigtable._generated import ( bigtable_instance_admin_pb2 as messages_v2_pb2) from tests.unit._testing import _FakeStub SERVE_NODES = 4 client = _Client(self.PROJECT) instance = _Instance(self.INSTANCE_ID, client) cluster = self._make_one( self.CLUSTER_ID, instance, serve_nodes=SERVE_NODES) # Create response_pb OP_ID = 5678 OP_NAME = ( 'operations/projects/%s/instances/%s/clusters/%s/operations/%d' % (self.PROJECT, self.INSTANCE_ID, self.CLUSTER_ID, OP_ID)) response_pb = operations_pb2.Operation(name=OP_NAME) # Patch the stub used by the API method. client._instance_stub = stub = _FakeStub(response_pb) # Perform the method and check the result. result = cluster.create() self.assertIsInstance(result, operation.Operation) self.assertEqual(result.operation.name, OP_NAME) self.assertIsNone(result.metadata) self.assertEqual(len(stub.method_calls), 1) api_name, args, kwargs = stub.method_calls[0] self.assertEqual(api_name, 'CreateCluster') request_pb, = args self.assertIsInstance(request_pb, messages_v2_pb2.CreateClusterRequest) self.assertEqual(request_pb.parent, instance.name) self.assertEqual(request_pb.cluster_id, self.CLUSTER_ID) self.assertEqual(request_pb.cluster.serve_nodes, SERVE_NODES) self.assertEqual(kwargs, {})
def test_reload(self): from tests.unit._testing import _FakeStub from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES SERVE_NODES = 31 LOCATION = 'LOCATION' client = _Client(self.PROJECT) instance = _Instance(self.INSTANCE_ID, client) cluster = self._make_one(self.CLUSTER_ID, instance) # Create request_pb request_pb = _GetClusterRequestPB(name=self.CLUSTER_NAME) # Create response_pb response_pb = _ClusterPB( serve_nodes=SERVE_NODES, location=LOCATION, ) # Patch the stub used by the API method. client._instance_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = None # reload() has no return value. # Check Cluster optional config values before. self.assertEqual(cluster.serve_nodes, DEFAULT_SERVE_NODES) # Perform the method and check the result. result = cluster.reload() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'GetCluster', (request_pb,), {}, )]) # Check Cluster optional config values before. self.assertEqual(cluster.serve_nodes, SERVE_NODES) self.assertEqual(cluster.location, LOCATION)
def test_do_mutate_retryable_rows_mismatch_num_responses(self): from google.cloud.bigtable.row import DirectRow from tests.unit._testing import _FakeStub client = _Client() instance = _Instance(self.INSTANCE_NAME, client=client) table = self._make_table(self.TABLE_ID, instance) row_1 = DirectRow(row_key=b'row_key', table=table) row_1.set_cell('cf', b'col', b'value1') row_2 = DirectRow(row_key=b'row_key_2', table=table) row_2.set_cell('cf', b'col', b'value2') response = self._make_responses([self.SUCCESS]) # Patch the stub used by the API method. client._data_stub = _FakeStub([response]) worker = self._make_worker( table._instance._client, table.name, [row_1, row_2]) with self.assertRaises(RuntimeError): worker._do_mutate_retryable_rows()
def test_create_w_explicit_serve_nodes(self): from google.api_core import operation from google.longrunning import operations_pb2 from google.cloud.bigtable._generated import ( bigtable_instance_admin_pb2 as messages_v2_pb2) from tests.unit._testing import _FakeStub SERVE_NODES = 5 client = _Client(self.PROJECT) instance = self._make_one(self.INSTANCE_ID, client, self.LOCATION_ID, serve_nodes=SERVE_NODES) # Create response_pb response_pb = operations_pb2.Operation(name=self.OP_NAME) # Patch the stub used by the API method. client._instance_stub = stub = _FakeStub(response_pb) # Perform the method and check the result. result = instance.create() self.assertIsInstance(result, operation.Operation) self.assertEqual(result.operation.name, self.OP_NAME) self.assertEqual(len(stub.method_calls), 1) api_name, args, kwargs = stub.method_calls[0] self.assertEqual(api_name, 'CreateInstance') request_pb, = args self.assertIsInstance(request_pb, messages_v2_pb2.CreateInstanceRequest) self.assertEqual(request_pb.parent, 'projects/%s' % (self.PROJECT,)) self.assertEqual(request_pb.instance_id, self.INSTANCE_ID) self.assertEqual(request_pb.instance.display_name, self.INSTANCE_ID) cluster = request_pb.clusters[self.INSTANCE_ID] self.assertEqual(cluster.serve_nodes, SERVE_NODES) self.assertEqual(kwargs, {})
def test_update(self): import datetime from google.api_core import operation from google.longrunning import operations_pb2 from google.protobuf.any_pb2 import Any from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.bigtable._generated import ( instance_pb2 as data_v2_pb2) from google.cloud.bigtable._generated import ( bigtable_instance_admin_pb2 as messages_v2_pb2) from tests.unit._testing import _FakeStub NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) SERVE_NODES = 81 client = _Client(self.PROJECT) instance = _Instance(self.INSTANCE_ID, client) cluster = self._make_one(self.CLUSTER_ID, instance, serve_nodes=SERVE_NODES) # Create request_pb request_pb = _ClusterPB( name=self.CLUSTER_NAME, serve_nodes=SERVE_NODES, ) # Create response_pb OP_ID = 5678 OP_NAME = ( 'operations/projects/%s/instances/%s/clusters/%s/operations/%d' % (self.PROJECT, self.INSTANCE_ID, self.CLUSTER_ID, OP_ID)) metadata = messages_v2_pb2.UpdateClusterMetadata(request_time=NOW_PB) type_url = 'type.googleapis.com/%s' % ( messages_v2_pb2.UpdateClusterMetadata.DESCRIPTOR.full_name,) response_pb = operations_pb2.Operation( name=OP_NAME, metadata=Any( type_url=type_url, value=metadata.SerializeToString() ) ) # Patch the stub used by the API method. client._instance_stub = stub = _FakeStub(response_pb) result = cluster.update() self.assertIsInstance(result, operation.Operation) self.assertEqual(result.operation.name, OP_NAME) self.assertIsInstance(result.metadata, messages_v2_pb2.UpdateClusterMetadata) self.assertEqual(result.metadata.request_time, NOW_PB) self.assertEqual(len(stub.method_calls), 1) api_name, args, kwargs = stub.method_calls[0] self.assertEqual(api_name, 'UpdateCluster') request_pb, = args self.assertIsInstance(request_pb, data_v2_pb2.Cluster) self.assertEqual(request_pb.name, self.CLUSTER_NAME) self.assertEqual(request_pb.serve_nodes, SERVE_NODES) self.assertEqual(kwargs, {})
def test_commit(self): from tests.unit._testing import _FakeStub from google.cloud.bigtable.row_filters import RowSampleFilter row_key = b'row_key' table_name = 'projects/more-stuff' column_family_id1 = u'column_family_id1' column_family_id2 = u'column_family_id2' column_family_id3 = u'column_family_id3' column1 = b'column1' column2 = b'column2' client = _Client() table = _Table(table_name, client=client) row_filter = RowSampleFilter(0.33) row = self._make_one(row_key, table, filter_=row_filter) # Create request_pb value1 = b'bytes-value' mutation1 = _MutationPB( set_cell=_MutationSetCellPB( family_name=column_family_id1, column_qualifier=column1, timestamp_micros=-1, # Default value. value=value1, ), ) mutation2 = _MutationPB( delete_from_row=_MutationDeleteFromRowPB(), ) mutation3 = _MutationPB( delete_from_column=_MutationDeleteFromColumnPB( family_name=column_family_id2, column_qualifier=column2, ), ) mutation4 = _MutationPB( delete_from_family=_MutationDeleteFromFamilyPB( family_name=column_family_id3, ), ) request_pb = _CheckAndMutateRowRequestPB( table_name=table_name, row_key=row_key, predicate_filter=row_filter.to_pb(), true_mutations=[mutation1, mutation3, mutation4], false_mutations=[mutation2], ) # Create response_pb predicate_matched = True response_pb = _CheckAndMutateRowResponsePB( predicate_matched=predicate_matched) # Patch the stub used by the API method. client._data_stub = stub = _FakeStub(response_pb) # Create expected_result. expected_result = predicate_matched # Perform the method and check the result. row.set_cell(column_family_id1, column1, value1, state=True) row.delete(state=False) row.delete_cell(column_family_id2, column2, state=True) row.delete_cells(column_family_id3, row.ALL_COLUMNS, state=True) result = row.commit() self.assertEqual(result, expected_result) self.assertEqual(stub.method_calls, [( 'CheckAndMutateRow', (request_pb,), {}, )]) self.assertEqual(row._true_pb_mutations, []) self.assertEqual(row._false_pb_mutations, [])