def test_do_mutate_retryable_rows_second_try_no_retryable(self): from google.cloud.bigtable.row import DirectRow from tests.unit._testing import _FakeStub # Setup: # - Mutate 2 rows. # - First try results: [success, non-retryable] # Action: # - Second try has no row to retry. # Expectation: # - After second try: [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') worker = self._make_worker(table._instance._client, table.name, [row_1, row_2]) worker.responses_statuses = self._make_responses_statuses( [self.SUCCESS, self.NON_RETRYABLE]) 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_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_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_do_mutate_retryable_rows_second_try_no_retryable(self): from google.cloud.bigtable.row import DirectRow # Setup: # - Mutate 2 rows. # - First try results: [success, non-retryable] # Action: # - Second try has no row to retry. # Expectation: # - After second try: [success, non-retryable] channel = self._make_channel() client = self._make_client(project='project-id', channel=channel, admin=True) instance = client.instance(instance_id=self.INSTANCE_ID) 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') worker = self._make_worker(client, table.name, [row_1, row_2]) worker.responses_statuses = self._make_responses_statuses( [self.SUCCESS, self.NON_RETRYABLE]) 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(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_second_try_no_retryable(self): from google.cloud.bigtable.row import DirectRow # Setup: # - Mutate 2 rows. # - First try results: [success, non-retryable] # Action: # - Second try has no row to retry. # Expectation: # - After second try: [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') worker = self._make_worker( table._instance._client, table.name, [row_1, row_2]) worker.responses_statuses = self._make_responses_statuses( [self.SUCCESS, self.NON_RETRYABLE]) 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(self): from google.cloud.bigtable.row import DirectRow # Setup: # - Mutate 2 rows. # Action: # - Initial attempt will mutate all 2 rows. # Expectation: # - Expect [success, non-retryable] channel = self._make_channel() client = self._make_client(project='project-id', channel=channel, admin=True) instance = client.instance(instance_id=self.INSTANCE_ID) 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. bigtable_stub = client._table_data_client.bigtable_stub bigtable_stub.MutateRows.side_effect = [[response]] worker = self._make_worker(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_callable_retry_timeout(self): from google.api_core.retry import Retry from google.cloud.bigtable._generated.bigtable_pb2 import ( MutateRowsResponse) from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.table import DEFAULT_RETRY from google.rpc.status_pb2 import Status # Setup: # - Mutate 2 rows. # Action: # - Initial attempt will mutate all 2 rows. # Expectation: # - Both rows always return retryable errors. # - google.api_core.Retry should keep retrying. # - Check MutateRows is called multiple times. # - By the time deadline is reached, statuses should be # [retryable, 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=4), ), MutateRowsResponse.Entry( index=1, status=Status(code=4), ), ], ) # Patch the stub used by the API method. client._data_stub = mock.MagicMock() client._data_stub.MutateRows.return_value = [response] retry = DEFAULT_RETRY.with_delay(initial=0.1, maximum=0.2, multiplier=2.0).with_deadline(0.5) worker = self._make_worker(client, table.name, [row_1, row_2]) statuses = worker(retry=retry) result = [status.code for status in statuses] expected_result = [4, 4] self.assertTrue(client._data_stub.MutateRows.call_count > 1) self.assertEqual(result, expected_result)
def test_mutate_row_with_max_mutations(self): table = _Table(self.TABLE_NAME) mutation_batcher = MutationsBatcher(table=table) row = DirectRow(row_key=b'row_key') row.set_cell('cf1', b'c1', 1) row.set_cell('cf1', b'c2', 2) row.set_cell('cf1', b'c3', 3) mutation_batcher.mutate(row) mutation_batcher.flush() self.assertEqual(table.mutation_calls, 1)
def test_mutation_batcher_mutate_w_max_mutations(): table = _Table(TABLE_NAME) mutation_batcher = _make_mutation_batcher(table=table) row = DirectRow(row_key=b"row_key") row.set_cell("cf1", b"c1", 1) row.set_cell("cf1", b"c2", 2) row.set_cell("cf1", b"c3", 3) mutation_batcher.mutate(row) mutation_batcher.flush() assert table.mutation_calls == 1
def test_mutate_row_with_max_mutations(self): table = _Table(self.TABLE_NAME) mutation_batcher = MutationsBatcher(table=table) row = DirectRow(row_key=b"row_key") row.set_cell("cf1", b"c1", 1) row.set_cell("cf1", b"c2", 2) row.set_cell("cf1", b"c3", 3) mutation_batcher.mutate(row) mutation_batcher.flush() self.assertEqual(table.mutation_calls, 1)
def test_mutate_row_with_max_mutations(self): table = _Table(self.TABLE_NAME) mutation_batcher = MutationsBatcher(table=table) row = DirectRow(row_key=b'row_key') row.set_cell('cf1', b'c1', 1) row.set_cell('cf1', b'c2', 2) row.set_cell('cf1', b'c3', 3) mutation_batcher.mutate(row) mutation_batcher.flush() self.assertEqual(table.mutation_calls, 1)
def test_mutate_row_with_max_mutations(self): table = _Table(self.TABLE_NAME) mutation_batcher = MutationsBatcher(table=table) row = DirectRow(row_key=b"row_key") row.set_cell("cf1", b"c1", 1) row.set_cell("cf1", b"c2", 2) row.set_cell("cf1", b"c3", 3) mutation_batcher.mutate(row) mutation_batcher.flush() self.assertEqual(table.mutation_calls, 1)
def test_mutate_row_with_max_mutations_failure(self): from google.cloud.bigtable.batcher import MaxMutationsError table = _Table(self.TABLE_NAME) mutation_batcher = MutationsBatcher(table=table) row = DirectRow(row_key=b'row_key') row.set_cell('cf1', b'c1', 1) row.set_cell('cf1', b'c2', 2) row.set_cell('cf1', b'c3', 3) row.set_cell('cf1', b'c4', 4) with self.assertRaises(MaxMutationsError): mutation_batcher.mutate(row)
def test_mutate_row_with_max_mutations_failure(self): from google.cloud.bigtable.batcher import MaxMutationsError table = _Table(self.TABLE_NAME) mutation_batcher = MutationsBatcher(table=table) row = DirectRow(row_key=b"row_key") row.set_cell("cf1", b"c1", 1) row.set_cell("cf1", b"c2", 2) row.set_cell("cf1", b"c3", 3) row.set_cell("cf1", b"c4", 4) with self.assertRaises(MaxMutationsError): mutation_batcher.mutate(row)
def test_mutate_row_with_max_mutations_failure(self): from google.cloud.bigtable.batcher import MaxMutationsError table = _Table(self.TABLE_NAME) mutation_batcher = MutationsBatcher(table=table) row = DirectRow(row_key=b'row_key') row.set_cell('cf1', b'c1', 1) row.set_cell('cf1', b'c2', 2) row.set_cell('cf1', b'c3', 3) row.set_cell('cf1', b'c4', 4) with self.assertRaises(MaxMutationsError): mutation_batcher.mutate(row)
def test_mutate_row_with_max_mutations_failure(self): from google.cloud.bigtable.batcher import MaxMutationsError table = _Table(self.TABLE_NAME) mutation_batcher = MutationsBatcher(table=table) row = DirectRow(row_key=b"row_key") row.set_cell("cf1", b"c1", 1) row.set_cell("cf1", b"c2", 2) row.set_cell("cf1", b"c3", 3) row.set_cell("cf1", b"c4", 4) with self.assertRaises(MaxMutationsError): mutation_batcher.mutate(row)
def test_mutate_row_with_max_row_bytes(self): table = _Table(self.TABLE_NAME) mutation_batcher = MutationsBatcher(table=table, max_row_bytes=3 * 1024 * 1024) number_of_bytes = 1 * 1024 * 1024 max_value = b"1" * number_of_bytes row = DirectRow(row_key=b"row_key") row.set_cell("cf1", b"c1", max_value) row.set_cell("cf1", b"c2", max_value) row.set_cell("cf1", b"c3", max_value) mutation_batcher.mutate(row) self.assertEqual(table.mutation_calls, 1)
def test_do_mutate_retryable_rows_second_retry(self): from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.table import _BigtableRetryableError # 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. channel = self._make_channel() client = self._make_client(project='project-id', channel=channel, admin=True) instance = client.instance(instance_id=self.INSTANCE_ID) 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. bigtable_stub = client._table_data_client.bigtable_stub bigtable_stub.MutateRows.side_effect = [[response]] worker = self._make_worker(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_mutate_row_with_max_row_bytes(self): table = _Table(self.TABLE_NAME) mutation_batcher = MutationsBatcher(table=table, max_row_bytes=3 * 1024 * 1024) number_of_bytes = 1 * 1024 * 1024 max_value = b'1' * number_of_bytes row = DirectRow(row_key=b'row_key') row.set_cell('cf1', b'c1', max_value) row.set_cell('cf1', b'c2', max_value) row.set_cell('cf1', b'c3', max_value) mutation_batcher.mutate(row) self.assertEqual(table.mutation_calls, 1)
def test_mutate_row_with_max_row_bytes(self): table = _Table(self.TABLE_NAME) mutation_batcher = MutationsBatcher(table=table, max_row_bytes=3 * 1024 * 1024) number_of_bytes = 1 * 1024 * 1024 max_value = b"1" * number_of_bytes row = DirectRow(row_key=b"row_key") row.set_cell("cf1", b"c1", max_value) row.set_cell("cf1", b"c2", max_value) row.set_cell("cf1", b"c3", max_value) mutation_batcher.mutate(row) self.assertEqual(table.mutation_calls, 1)
def test_mutate_row_with_max_row_bytes(self): table = _Table(self.TABLE_NAME) mutation_batcher = MutationsBatcher(table=table, max_row_bytes=3 * 1024 * 1024) number_of_bytes = 1 * 1024 * 1024 max_value = b'1' * number_of_bytes row = DirectRow(row_key=b'row_key') row.set_cell('cf1', b'c1', max_value) row.set_cell('cf1', b'c2', max_value) row.set_cell('cf1', b'c3', max_value) mutation_batcher.mutate(row) self.assertEqual(table.mutation_calls, 1)
def test_mutation_batcher_mutate_w_max_row_bytes(): table = _Table(TABLE_NAME) mutation_batcher = _make_mutation_batcher(table=table, max_row_bytes=3 * 1024 * 1024) number_of_bytes = 1 * 1024 * 1024 max_value = b"1" * number_of_bytes row = DirectRow(row_key=b"row_key") row.set_cell("cf1", b"c1", max_value) row.set_cell("cf1", b"c2", max_value) row.set_cell("cf1", b"c3", max_value) mutation_batcher.mutate(row) assert table.mutation_calls == 1
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_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_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_callable_retry_timeout(self): from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.table import DEFAULT_RETRY # Setup: # - Mutate 2 rows. # Action: # - Initial attempt will mutate all 2 rows. # Expectation: # - Both rows always return retryable errors. # - google.api_core.Retry should keep retrying. # - Check MutateRows is called multiple times. # - By the time deadline is reached, statuses should be # [retryable, retryable] channel = self._make_channel() client = self._make_client(project='project-id', channel=channel, admin=True) instance = client.instance(instance_id=self.INSTANCE_ID) 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.RETRYABLE_1, self.RETRYABLE_1]) # Patch the stub used by the API method. bigtable_stub = client._table_data_client.bigtable_stub bigtable_stub.MutateRows.return_value = [response] retry = DEFAULT_RETRY.with_delay( initial=0.1, maximum=0.2, multiplier=2.0).with_deadline(0.5) worker = self._make_worker(client, table.name, [row_1, row_2]) statuses = worker(retry=retry) result = [status.code for status in statuses] expected_result = [self.RETRYABLE_1, self.RETRYABLE_1] self.assertTrue( client._table_data_client.bigtable_stub.MutateRows.call_count > 1) self.assertEqual(result, expected_result)
def test_do_mutate_retryable_rows_second_try(self): from google.cloud.bigtable.row import DirectRow # 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] channel = self._make_channel() client = self._make_client(project='project-id', channel=channel, admin=True) instance = client.instance(instance_id=self.INSTANCE_ID) 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. bigtable_stub = client._table_data_client.bigtable_stub bigtable_stub.MutateRows.side_effect = [[response]] worker = self._make_worker(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_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_callable_retry_timeout(self): from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.table import DEFAULT_RETRY # Setup: # - Mutate 2 rows. # Action: # - Initial attempt will mutate all 2 rows. # Expectation: # - Both rows always return retryable errors. # - google.api_core.Retry should keep retrying. # - Check MutateRows is called multiple times. # - By the time deadline is reached, statuses should be # [retryable, 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.RETRYABLE_1, self.RETRYABLE_1]) # Patch the stub used by the API method. client._data_stub = mock.MagicMock() client._data_stub.MutateRows.return_value = [response] retry = DEFAULT_RETRY.with_delay( initial=0.1, maximum=0.2, multiplier=2.0).with_deadline(0.5) worker = self._make_worker(client, table.name, [row_1, row_2]) statuses = worker(retry=retry) result = [status.code for status in statuses] expected_result = [self.RETRYABLE_1, self.RETRYABLE_1] self.assertTrue(client._data_stub.MutateRows.call_count > 1) self.assertEqual(result, expected_result)
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_do_mutate_retryable_rows_mismatch_num_responses(self): from google.cloud.bigtable.row import DirectRow channel = self._make_channel() client = self._make_client(project='project-id', channel=channel, admin=True) instance = client.instance(instance_id=self.INSTANCE_ID) 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. bigtable_stub = client._table_data_client.bigtable_stub bigtable_stub.MutateRows.side_effect = [[response]] worker = self._make_worker(client, table.name, [row_1, row_2]) with self.assertRaises(RuntimeError): worker._do_mutate_retryable_rows()
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_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_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_callable_retry(self): from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.table import DEFAULT_RETRY # Setup: # - Mutate 3 rows. # Action: # - Initial attempt will mutate all 3 rows. # Expectation: # - First attempt will result in one retryable error. # - Second attempt will result in success for the retry-ed row. # - Check MutateRows is called twice. # - State of responses_statuses should be # [success, 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') row_3 = DirectRow(row_key=b'row_key_3', table=table) row_3.set_cell('cf', b'col', b'value3') response_1 = self._make_responses([ self.SUCCESS, self.RETRYABLE_1, self.NON_RETRYABLE]) response_2 = self._make_responses([self.SUCCESS]) # Patch the stub used by the API method. client._data_stub = mock.MagicMock() client._data_stub.MutateRows.side_effect = [[response_1], [response_2]] retry = DEFAULT_RETRY.with_delay(initial=0.1) worker = self._make_worker(client, table.name, [row_1, row_2, row_3]) statuses = worker(retry=retry) result = [status.code for status in statuses] expected_result = [self.SUCCESS, self.SUCCESS, self.NON_RETRYABLE] client._data_stub.MutateRows.assert_has_calls([ mock.call(mock.ANY), mock.call(mock.ANY)]) self.assertEqual(client._data_stub.MutateRows.call_count, 2) self.assertEqual(result, expected_result)
def test_callable_retry(self): from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.table import DEFAULT_RETRY # Setup: # - Mutate 3 rows. # Action: # - Initial attempt will mutate all 3 rows. # Expectation: # - First attempt will result in one retryable error. # - Second attempt will result in success for the retry-ed row. # - Check MutateRows is called twice. # - State of responses_statuses should be # [success, success, non-retryable] channel = self._make_channel() client = self._make_client(project='project-id', channel=channel, admin=True) instance = client.instance(instance_id=self.INSTANCE_ID) 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_1 = self._make_responses([ self.SUCCESS, self.RETRYABLE_1, self.NON_RETRYABLE]) response_2 = self._make_responses([self.SUCCESS]) # Patch the stub used by the API method. client._table_data_client.bigtable_stub.MutateRows.side_effect = [ [response_1], [response_2]] retry = DEFAULT_RETRY.with_delay(initial=0.1) worker = self._make_worker(client, table.name, [row_1, row_2, row_3]) statuses = worker(retry=retry) result = [status.code for status in statuses] expected_result = [self.SUCCESS, self.SUCCESS, self.NON_RETRYABLE] self.assertEqual( client._table_data_client.bigtable_stub.MutateRows.call_count, 2) self.assertEqual(result, expected_result)
def test_do_mutate_retryable_rows_retry(self): from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.table import _BigtableRetryableError # 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] channel = self._make_channel() client = self._make_client(project='project-id', channel=channel, admin=True) instance = client.instance(instance_id=self.INSTANCE_ID) 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. bigtable_stub = client._table_data_client.bigtable_stub bigtable_stub.MutateRows.side_effect = [[response]] worker = self._make_worker(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_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_callable_no_retry_strategy(self): from google.cloud.bigtable.row import DirectRow # Setup: # - Mutate 3 rows. # Action: # - Attempt to mutate the rows w/o any retry strategy. # Expectation: # - Since no retry, should return statuses as they come back. # - Even if there are retryable errors, no retry attempt is made. # - State of responses_statuses should be # [success, retryable, non-retryable] channel = self._make_channel() client = self._make_client(project='project-id', channel=channel, admin=True) instance = client.instance(instance_id=self.INSTANCE_ID) 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. bigtable_stub = client._table_data_client.bigtable_stub bigtable_stub.MutateRows.return_value = [response] worker = self._make_worker(client, table.name, [row_1, row_2, row_3]) statuses = worker(retry=None) result = [status.code for status in statuses] expected_result = [self.SUCCESS, self.RETRYABLE_1, self.NON_RETRYABLE] client._table_data_client.bigtable_stub.MutateRows.assert_called_once() self.assertEqual(result, expected_result)
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_callable_no_retry_strategy(self): from google.cloud.bigtable.row import DirectRow # Setup: # - Mutate 3 rows. # Action: # - Attempt to mutate the rows w/o any retry strategy. # Expectation: # - Since no retry, should return statuses as they come back. # - Even if there are retryable errors, no retry attempt is made. # - 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 = mock.MagicMock() client._data_stub.MutateRows.return_value = [response] worker = self._make_worker(client, table.name, [row_1, row_2, row_3]) statuses = worker(retry=None) result = [status.code for status in statuses] expected_result = [self.SUCCESS, self.RETRYABLE_1, self.NON_RETRYABLE] client._data_stub.MutateRows.assert_called_once() self.assertEqual(result, expected_result)
def test_callable_retry(self): from google.api_core.retry import Retry from google.cloud.bigtable._generated.bigtable_pb2 import ( MutateRowsResponse) from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.table import DEFAULT_RETRY from google.rpc.status_pb2 import Status # Setup: # - Mutate 3 rows. # Action: # - Initial attempt will mutate all 3 rows. # Expectation: # - First attempt will result in one retryable error. # - Second attempt will result in success for the retry-ed row. # - Check MutateRows is called twice. # - State of responses_statuses should be # [success, 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') row_3 = DirectRow(row_key=b'row_key_3', table=table) row_3.set_cell('cf', b'col', b'value3') response_1 = 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), ), ], ) response_2 = MutateRowsResponse(entries=[ MutateRowsResponse.Entry( index=0, status=Status(code=0), ), ], ) # Patch the stub used by the API method. client._data_stub = mock.MagicMock() client._data_stub.MutateRows.side_effect = [[response_1], [response_2]] retry = DEFAULT_RETRY.with_delay(initial=0.1) worker = self._make_worker(client, table.name, [row_1, row_2, row_3]) statuses = worker(retry=retry) result = [status.code for status in statuses] expected_result = [0, 0, 1] client._data_stub.MutateRows.assert_has_calls( [mock.call(mock.ANY), mock.call(mock.ANY)]) self.assertEqual(client._data_stub.MutateRows.call_count, 2) self.assertEqual(result, expected_result)
def test_do_mutate_retryable_rows_second_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 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 = MutateRowsResponse(entries=[ MutateRowsResponse.Entry( index=0, status=Status(code=0), ), MutateRowsResponse.Entry( index=1, status=Status(code=4), ), ], ) # 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]) with self.assertRaises(ServiceUnavailable): worker._do_mutate_retryable_rows() statuses = worker.responses_statuses result = [status.code for status in statuses] expected_result = [0, 0, 1, 4] self.assertEqual(result, expected_result)