implemented full success tests for rpcs

Author: daniel-sanche

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

@@ -1108,9 +1108,7 @@ async def read_rows(
         )
         return [row async for row in row_generator]
 
-    @CrossSync.convert(
-        replace_symbols={"__anext__": "__next__", "StopAsyncIteration": "StopIteration"}
-    )
+    @CrossSync.convert
     async def read_row(
         self,
         row_key: str | bytes,
@@ -1168,8 +1166,9 @@ async def read_row(
         )
         results_generator = row_merger.start_operation()
         try:
-            return await results_generator.__anext__()
-        except StopAsyncIteration:
+            results = [a async for a in results_generator]
+            return results[0]
+        except IndexError:
             return None
 
     @CrossSync.convert

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, CompletedAttemptMetric
+from google.cloud.bigtable.data._metrics.data_model import CompletedOperationMetric, CompletedAttemptMetric, ActiveOperationMetric, OperationState
 
 from google.cloud.bigtable.data._cross_sync import CrossSync
 
@@ -96,42 +96,275 @@ async def temp_rows(self, target):
     @CrossSync.convert
     @CrossSync.pytest_fixture(scope="session")
     async def target(self, client, table_id, instance_id, handler):
-        """
-        This fixture runs twice: once for a standard table, and once with an authorized view
-
-        Note: emulator doesn't support authorized views. Only use target
-        """
         async with client.get_table(instance_id, table_id) as table:
             table._metrics.add_handler(handler)
             yield table
 
     @CrossSync.pytest
     async def test_read_rows(self, target, temp_rows, handler, cluster_config):
-        pass
+        await temp_rows.add_row(b"row_key_1")
+        await temp_rows.add_row(b"row_key_2")
+        handler.clear()
+        row_list = await target.read_rows({})
+        assert len(row_list) == 2
+        # 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.is_streaming is True
+        assert operation.op_type.value == "ReadRows"
+        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 not None and operation.first_response_latency_ns < operation.duration_ns
+        assert operation.flow_throttling_time_ns == 0
+        # validate attempt
+        attempt = handler.completed_attempts[0]
+        assert isinstance(attempt, CompletedAttemptMetric)
+        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 and attempt.application_blocking_time_ns < operation.duration_ns
+        assert attempt.grpc_throttling_time_ns == 0  # TODO: confirm
 
     @CrossSync.pytest
     async def test_read_rows_stream(self, target, temp_rows, handler, cluster_config):
-        pass
+        await temp_rows.add_row(b"row_key_1")
+        await temp_rows.add_row(b"row_key_2")
+        handler.clear()
+        # full table scan
+        generator = await target.read_rows_stream({})
+        row_list = [r async for r in generator]
+        assert len(row_list) == 2
+        # 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.is_streaming is True
+        assert operation.op_type.value == "ReadRows"
+        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 not None and operation.first_response_latency_ns < operation.duration_ns
+        assert operation.flow_throttling_time_ns == 0
+        # validate attempt
+        attempt = handler.completed_attempts[0]
+        assert isinstance(attempt, CompletedAttemptMetric)
+        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 and attempt.application_blocking_time_ns < operation.duration_ns
+        assert attempt.grpc_throttling_time_ns == 0  # TODO: confirm
 
     @CrossSync.pytest
     async def test_read_row(self, target, temp_rows, handler, cluster_config):
-        pass
+        await temp_rows.add_row(b"row_key_1")
+        handler.clear()
+        await target.read_row(b"row_key_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.is_streaming is False
+        assert operation.op_type.value == "ReadRows"
+        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 > 0 and operation.first_response_latency_ns < operation.duration_ns
+        assert operation.flow_throttling_time_ns == 0
+        # validate attempt
+        attempt = handler.completed_attempts[0]
+        assert isinstance(attempt, CompletedAttemptMetric)
+        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 and attempt.application_blocking_time_ns < operation.duration_ns
+        assert attempt.grpc_throttling_time_ns == 0  # TODO: confirm
 
     @CrossSync.pytest
     async def test_read_rows_sharded(self, target, temp_rows, handler, cluster_config):
-        pass
+        from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
+        await temp_rows.add_row(b"a")
+        await temp_rows.add_row(b"b")
+        await temp_rows.add_row(b"c")
+        await temp_rows.add_row(b"d")
+        query1 = ReadRowsQuery(row_keys=[b"a", b"c"])
+        query2 = ReadRowsQuery(row_keys=[b"b", b"d"])
+        handler.clear()
+        row_list = await target.read_rows_sharded([query1, query2])
+        assert len(row_list) == 4
+        # validate counts
+        assert len(handler.completed_operations) == 2
+        assert len(handler.completed_attempts) == 2
+        assert len(handler.cancelled_operations) == 0
+        # validate operations
+        for operation in handler.completed_operations:
+            assert isinstance(operation, CompletedOperationMetric)
+            assert operation.final_status.value[0] == 0
+            assert operation.is_streaming is True
+            assert operation.op_type.value == "ReadRows"
+            assert len(operation.completed_attempts) == 1
+            attempt = operation.completed_attempts[0]
+            assert attempt in handler.completed_attempts
+            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 not None and operation.first_response_latency_ns < operation.duration_ns
+            assert operation.flow_throttling_time_ns == 0
+            # validate attempt
+            assert isinstance(attempt, CompletedAttemptMetric)
+            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 and attempt.application_blocking_time_ns < operation.duration_ns
+            assert attempt.grpc_throttling_time_ns == 0  # TODO: confirm
 
     @CrossSync.pytest
     async def test_bulk_mutate_rows(self, target, temp_rows, handler, cluster_config):
-        pass
+        from google.cloud.bigtable.data.mutations import RowMutationEntry
+
+        new_value = uuid.uuid4().hex.encode()
+        row_key, mutation = await temp_rows.create_row_and_mutation(
+            target, new_value=new_value
+        )
+        bulk_mutation = RowMutationEntry(row_key, [mutation])
+
+        handler.clear()
+        await target.bulk_mutate_rows([bulk_mutation])
+        # 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.is_streaming is False
+        assert operation.op_type.value == "MutateRows"
+        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 isinstance(attempt, CompletedAttemptMetric)
+        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
 
     @CrossSync.pytest
-    async def test_row_batcher(self, target, temp_rows, handler, cluster_config):
-        pass
+    async def test_mutate_rows_batcher(self, target, temp_rows, handler, cluster_config):
+        from google.cloud.bigtable.data.mutations import RowMutationEntry
+
+        new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)]
+        row_key, mutation = await temp_rows.create_row_and_mutation(
+            target, new_value=new_value
+        )
+        row_key2, mutation2 = await temp_rows.create_row_and_mutation(
+            target, new_value=new_value2
+        )
+        bulk_mutation = RowMutationEntry(row_key, [mutation])
+        bulk_mutation2 = RowMutationEntry(row_key2, [mutation2])
+
+        handler.clear()
+        async with target.mutations_batcher() as batcher:
+            await batcher.append(bulk_mutation)
+            await batcher.append(bulk_mutation2)
+        # validate counts
+        assert len(handler.completed_operations) == 1
+        assert len(handler.completed_attempts) == 1
+        # bacher expects to cancel staged operation on close
+        assert len(handler.cancelled_operations) == 1
+        cancelled = handler.cancelled_operations[0]
+        assert isinstance(cancelled, ActiveOperationMetric)
+        assert cancelled.state == OperationState.CREATED
+        assert not cancelled.completed_attempts
+        # validate operation
+        operation = handler.completed_operations[0]
+        assert isinstance(operation, CompletedOperationMetric)
+        assert operation.final_status.value[0] == 0
+        assert operation.is_streaming is False
+        assert operation.op_type.value == "MutateRows"
+        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 and operation.flow_throttling_time_ns < operation.duration_ns
+        # validate attempt
+        attempt = handler.completed_attempts[0]
+        assert isinstance(attempt, CompletedAttemptMetric)
+        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
 
     @CrossSync.pytest
     async def test_mutate_row(self, target, temp_rows, handler, cluster_config):
-        pass
+        row_key = b"bulk_mutate"
+        new_value = uuid.uuid4().hex.encode()
+        row_key, mutation = await temp_rows.create_row_and_mutation(
+            target, new_value=new_value
+        )
+        handler.clear()
+        await target.mutate_row(row_key, mutation)
+        # 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.is_streaming is False
+        assert operation.op_type.value == "MutateRow"
+        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 isinstance(attempt, CompletedAttemptMetric)
+        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
 
     @CrossSync.pytest
     async def test_sample_row_keys(self, target, temp_rows, handler, cluster_config):
@@ -155,6 +388,7 @@ async def test_sample_row_keys(self, target, temp_rows, handler, cluster_config)
         assert operation.flow_throttling_time_ns == 0
         # validate attempt
         attempt = handler.completed_attempts[0]
+        assert isinstance(attempt, CompletedAttemptMetric)
         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
@@ -193,6 +427,7 @@ async def test_read_modify_write(self, target, temp_rows, handler, cluster_confi
         assert operation.flow_throttling_time_ns == 0
         # validate attempt
         attempt = handler.completed_attempts[0]
+        assert isinstance(attempt, CompletedAttemptMetric)
         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
@@ -247,6 +482,7 @@ async def test_check_and_mutate_row(self, target, temp_rows, handler, cluster_co
         assert operation.flow_throttling_time_ns == 0
         # validate attempt
         attempt = handler.completed_attempts[0]
+        assert isinstance(attempt, CompletedAttemptMetric)
         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