added test for read_modify_write

Author: daniel-sanche

google/cloud/bigtable/data/_async/client.py

@@ -1666,7 +1666,7 @@ async def read_modify_write_row(
         if not rules:
             raise ValueError("rules must contain at least one item")
 
-        with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE):
+        with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE) as op:
             result = await self.client._gapic_client.read_modify_write_row(
                 request=ReadModifyWriteRowRequest(
                     rules=[rule._to_pb() for rule in rules],
@@ -1678,6 +1678,7 @@ async def read_modify_write_row(
                 ),
                 timeout=operation_timeout,
                 retry=None,
+                metadata=[op.interceptor_metadata],
             )
             # construct Row from result
             return Row._from_pb(result.row)

tests/system/data/test_metrics_async.py

@@ -17,7 +17,7 @@
 import uuid
 
 from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler
-from google.cloud.bigtable.data._metrics.data_model import CompletedOperationMetric
+from google.cloud.bigtable.data._metrics.data_model import CompletedOperationMetric, CompletedAttemptMetric
 
 from google.cloud.bigtable.data._cross_sync import CrossSync
 
@@ -43,11 +43,8 @@ def on_operation_complete(self, op):
     def on_operation_cancelled(self, op):
         self.cancelled_operations.append(op)
 
-    def on_attempt_complete(self, attempt, op):
-        self.completed_attempts.append((attempt, op))
-
-    def total(self):
-        return len(self.completed_operations) + len(self._ancelled_operations) + len(self.completed_attempts)
+    def on_attempt_complete(self, attempt, _):
+        self.completed_attempts.append(attempt)
 
     def clear(self):
         self.cancelled_operations.clear()
@@ -101,7 +98,7 @@ async def target(self, client, table_id, instance_id, handler):
             yield table
 
     @CrossSync.pytest
-    async def test_read_modify_write(self, target, temp_rows, handler):
+    async def test_read_modify_write(self, target, temp_rows, handler, cluster_config):
         from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule
 
         row_key = b"test-row-key"
@@ -112,11 +109,28 @@ async def test_read_modify_write(self, target, temp_rows, handler):
         )
         rule = IncrementRule(family, qualifier, 1)
         await target.read_modify_write_row(row_key, rule)
-        breakpoint()
-        assert handler.total() == 1
+        # validate counts
         assert len(handler.completed_operations) == 1
         assert len(handler.completed_attempts) == 1
         assert len(handler.cancelled_operations) == 0
+        # validate operation
         operation = handler.completed_operations[0]
+        assert isinstance(operation, CompletedOperationMetric)
         assert operation.final_status.value[0] == 0
-        assert operation.final_status.value[0] == 0
+        assert operation.is_streaming is False
+        assert operation.op_type.value == "ReadModifyWriteRow"
+        assert len(operation.completed_attempts) == 1
+        assert operation.completed_attempts[0] == handler.completed_attempts[0]
+        assert operation.cluster_id == next(iter(cluster_config.keys()))
+        assert operation.zone == cluster_config[operation.cluster_id].location.split("/")[-1]
+        assert operation.duration_ns > 0 and operation.duration_ns < 1e9
+        assert operation.first_response_latency_ns is None  # populated for read_rows only
+        assert operation.flow_throttling_time_ns == 0
+        # validate attempt
+        attempt = handler.completed_attempts[0]
+        assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns
+        assert attempt.end_status.value[0] == 0
+        assert attempt.backoff_before_attempt_ns == 0
+        assert attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns
+        assert attempt.application_blocking_time_ns == 0
+        assert attempt.grpc_throttling_time_ns == 0  # TODO: confirm