def test_write_metrics(self):
MetricsEnvironment.process_wide_container().reset()
write_fn = bigtableio._BigTableWriteFn(self._PROJECT_ID,
self._INSTANCE_ID,
self._TABLE_ID)
write_fn.table = self.table
write_fn.start_bundle()
number_of_rows = 2
error = Status()
error.message = 'Entity already exists.'
error.code = ALREADY_EXISTS
success = Status()
success.message = 'Success'
success.code = OK
rows_response = [error, success] * number_of_rows
with patch.object(Table, 'mutate_rows', return_value=rows_response):
direct_rows = [
self.generate_row(i) for i in range(number_of_rows * 2)
]
for direct_row in direct_rows:
write_fn.process(direct_row)
try:
write_fn.finish_bundle()
except: # pylint: disable=bare-except
# Currently we fail the bundle when there are any failures.
# TODO(BEAM-13849): remove after bigtableio can selectively retry.
pass
self.verify_write_call_metric(
self._PROJECT_ID, self._INSTANCE_ID, self._TABLE_ID,
ServiceCallMetric.bigtable_error_code_to_grpc_status_string(
ALREADY_EXISTS), 2)
self.verify_write_call_metric(
self._PROJECT_ID, self._INSTANCE_ID, self._TABLE_ID,
ServiceCallMetric.bigtable_error_code_to_grpc_status_string(
OK), 2)
def test_write_metrics(self):
MetricsEnvironment.process_wide_container().reset()
write_fn = bigtableio._BigTableWriteFn(self._PROJECT_ID,
self._INSTANCE_ID,
self._TABLE_ID)
write_fn.table = self.table
write_fn.start_bundle()
number_of_rows = 2
error = Status()
error.message = 'Entity already exists.'
error.code = ALREADY_EXISTS
success = Status()
success.message = 'Success'
success.code = OK
rows_response = [error, success] * number_of_rows
with patch.object(Table, 'mutate_rows', return_value=rows_response):
direct_rows = [
self.generate_row(i) for i in range(number_of_rows * 2)
]
for direct_row in direct_rows:
write_fn.process(direct_row)
write_fn.finish_bundle()
self.verify_write_call_metric(
self._PROJECT_ID, self._INSTANCE_ID, self._TABLE_ID,
ServiceCallMetric.bigtable_error_code_to_grpc_status_string(
ALREADY_EXISTS), 2)
self.verify_write_call_metric(
self._PROJECT_ID, self._INSTANCE_ID, self._TABLE_ID,
ServiceCallMetric.bigtable_error_code_to_grpc_status_string(
OK), 2)