fixed more tests

Author: daniel-sanche

tests/unit/data/_async/test__mutate_rows.py

@@ -17,6 +17,7 @@
 from google.cloud.bigtable_v2.types import MutateRowsResponse
 from google.cloud.bigtable.data.mutations import RowMutationEntry
 from google.cloud.bigtable.data.mutations import DeleteAllFromRow
+from google.cloud.bigtable.data._metrics import ActiveOperationMetric
 from google.rpc import status_pb2
 from google.api_core.exceptions import DeadlineExceeded
 from google.api_core.exceptions import Forbidden
@@ -48,6 +49,7 @@ def _make_one(self, *args, **kwargs):
             kwargs["attempt_timeout"] = kwargs.pop("attempt_timeout", 0.1)
             kwargs["retryable_exceptions"] = kwargs.pop("retryable_exceptions", ())
             kwargs["mutation_entries"] = kwargs.pop("mutation_entries", [])
+            kwargs["metric"] = kwargs.pop("metric", ActiveOperationMetric("MUTATE_ROWS"))
         return self._target_class()(*args, **kwargs)
 
     def _make_mutation(self, count=1, size=1):
@@ -90,13 +92,15 @@ def test_ctor(self):
         entries = [self._make_mutation(), self._make_mutation()]
         operation_timeout = 0.05
         attempt_timeout = 0.01
+        metric = mock.Mock()
         retryable_exceptions = ()
         instance = self._make_one(
             client,
             table,
             entries,
             operation_timeout,
             attempt_timeout,
+            metric,
             retryable_exceptions,
         )
         # running gapic_fn should trigger a client call with baked-in args
@@ -116,6 +120,7 @@ def test_ctor(self):
         assert instance.is_retryable(RuntimeError("")) is False
         assert instance.remaining_indices == list(range(len(entries)))
         assert instance.errors == {}
+        assert instance._operation_metric == metric
 
     def test_ctor_too_many_entries(self):
         """
@@ -139,6 +144,7 @@ def test_ctor_too_many_entries(self):
                 entries,
                 operation_timeout,
                 attempt_timeout,
+                mock.Mock(),
             )
         assert "mutate_rows requests can contain at most 100000 mutations" in str(
             e.value
@@ -152,14 +158,15 @@ async def test_mutate_rows_operation(self):
         """
         client = mock.Mock()
         table = mock.Mock()
+        metric = ActiveOperationMetric("MUTATE_ROWS")
         entries = [self._make_mutation(), self._make_mutation()]
         operation_timeout = 0.05
         cls = self._target_class()
         with mock.patch(
             f"{cls.__module__}.{cls.__name__}._run_attempt", CrossSync.Mock()
         ) as attempt_mock:
             instance = self._make_one(
-                client, table, entries, operation_timeout, operation_timeout
+                client, table, entries, operation_timeout, operation_timeout, metric
             )
             await instance.start()
             assert attempt_mock.call_count == 1
@@ -173,6 +180,7 @@ async def test_mutate_rows_attempt_exception(self, exc_type):
         client = CrossSync.Mock()
         table = mock.Mock()
         table._request_path = {"table_name": "table"}
+        metric = ActiveOperationMetric("MUTATE_ROWS")
         table.app_profile_id = None
         entries = [self._make_mutation(), self._make_mutation()]
         operation_timeout = 0.05
@@ -181,7 +189,7 @@ async def test_mutate_rows_attempt_exception(self, exc_type):
         found_exc = None
         try:
             instance = self._make_one(
-                client, table, entries, operation_timeout, operation_timeout
+                client, table, entries, operation_timeout, operation_timeout, metric
             )
             await instance._run_attempt()
         except Exception as e:
@@ -203,6 +211,7 @@ async def test_mutate_rows_exception(self, exc_type):
 
         client = mock.Mock()
         table = mock.Mock()
+        metric = ActiveOperationMetric("MUTATE_ROWS")
         entries = [self._make_mutation(), self._make_mutation()]
         operation_timeout = 0.05
         expected_cause = exc_type("abort")
@@ -215,7 +224,7 @@ async def test_mutate_rows_exception(self, exc_type):
             found_exc = None
             try:
                 instance = self._make_one(
-                    client, table, entries, operation_timeout, operation_timeout
+                    client, table, entries, operation_timeout, operation_timeout, metric
                 )
                 await instance.start()
             except MutationsExceptionGroup as e:
@@ -239,6 +248,7 @@ async def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type):
 
         client = mock.Mock()
         table = mock.Mock()
+        metric = ActiveOperationMetric("MUTATE_ROWS")
         entries = [self._make_mutation()]
         operation_timeout = 1
         expected_cause = exc_type("retry")
@@ -255,6 +265,7 @@ async def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type):
                 entries,
                 operation_timeout,
                 operation_timeout,
+                metric,
                 retryable_exceptions=(exc_type,),
             )
             await instance.start()
@@ -271,6 +282,7 @@ async def test_mutate_rows_incomplete_ignored(self):
 
         client = mock.Mock()
         table = mock.Mock()
+        metric = ActiveOperationMetric("MUTATE_ROWS")
         entries = [self._make_mutation()]
         operation_timeout = 0.05
         with mock.patch.object(
@@ -282,7 +294,7 @@ async def test_mutate_rows_incomplete_ignored(self):
             found_exc = None
             try:
                 instance = self._make_one(
-                    client, table, entries, operation_timeout, operation_timeout
+                    client, table, entries, operation_timeout, operation_timeout, metric
                 )
                 await instance.start()
             except MutationsExceptionGroup as e:

tests/unit/data/_async/test__read_rows.py

@@ -15,6 +15,7 @@
 import pytest
 
 from google.cloud.bigtable.data._cross_sync import CrossSync
+from google.cloud.bigtable.data._metrics import ActiveOperationMetric
 
 # try/except added for compatibility with python < 3.8
 try:
@@ -59,6 +60,7 @@ def test_ctor(self):
         expected_operation_timeout = 42
         expected_request_timeout = 44
         time_gen_mock = mock.Mock()
+        expected_metric = mock.Mock()
         subpath = "_async" if CrossSync.is_async else "_sync_autogen"
         with mock.patch(
             f"google.cloud.bigtable.data.{subpath}._read_rows._attempt_timeout_generator",
@@ -69,6 +71,7 @@ def test_ctor(self):
                 table,
                 operation_timeout=expected_operation_timeout,
                 attempt_timeout=expected_request_timeout,
+                metric=expected_metric,
             )
         assert time_gen_mock.call_count == 1
         time_gen_mock.assert_called_once_with(
@@ -81,6 +84,7 @@ def test_ctor(self):
         assert instance.request.table_name == "test_table"
         assert instance.request.app_profile_id == table.app_profile_id
         assert instance.request.rows_limit == row_limit
+        assert instance._operation_metric == expected_metric
 
     @pytest.mark.parametrize(
         "in_keys,last_key,expected",
@@ -269,7 +273,7 @@ async def mock_stream():
         table = mock.Mock()
         table._request_path = {"table_name": "table_name"}
         table.app_profile_id = "app_profile_id"
-        instance = self._make_one(query, table, 10, 10)
+        instance = self._make_one(query, table, 10, 10, ActiveOperationMetric("READ_ROWS"))
         assert instance._remaining_count == start_limit
         # read emit_num rows
         async for val in instance.chunk_stream(awaitable_stream()):
@@ -308,7 +312,7 @@ async def mock_stream():
         table = mock.Mock()
         table._request_path = {"table_name": "table_name"}
         table.app_profile_id = "app_profile_id"
-        instance = self._make_one(query, table, 10, 10)
+        instance = self._make_one(query, table, 10, 10, ActiveOperationMetric("READ_ROWS"))
         assert instance._remaining_count == start_limit
         with pytest.raises(InvalidChunk) as e:
             # read emit_num rows
@@ -334,7 +338,7 @@ async def mock_stream():
         with mock.patch.object(
             self._get_target_class(), "_read_rows_attempt"
         ) as mock_attempt:
-            instance = self._make_one(mock.Mock(), mock.Mock(), 1, 1)
+            instance = self._make_one(mock.Mock(), mock.Mock(), 1, 1, ActiveOperationMetric("READ_ROWS"))
             wrapped_gen = mock_stream()
             mock_attempt.return_value = wrapped_gen
             gen = instance.start_operation()

tests/unit/data/_async/test_client.py

@@ -1205,7 +1205,6 @@ async def test_ctor(self):
         assert instance_key in client._active_instances
         assert client._instance_owners[instance_key] == {id(table)}
         assert isinstance(table._metrics, BigtableClientSideMetricsController)
-        assert table._metrics.interceptor == client._metrics_interceptor
         assert table.default_operation_timeout == expected_operation_timeout
         assert table.default_attempt_timeout == expected_attempt_timeout
         assert (
@@ -1547,7 +1546,6 @@ async def test_ctor(self):
         assert instance_key in client._active_instances
         assert client._instance_owners[instance_key] == {id(view)}
         assert isinstance(view._metrics, BigtableClientSideMetricsController)
-        assert view._metrics.interceptor == client._metrics_interceptor
         assert view.default_operation_timeout == expected_operation_timeout
         assert view.default_attempt_timeout == expected_attempt_timeout
         assert (
@@ -1959,9 +1957,19 @@ async def test_read_row(self):
         async with self._make_client() as client:
             table = client.get_table("instance", "table")
             row_key = b"test_1"
-            with mock.patch.object(table, "read_rows") as read_rows:
+            with mock.patch.object(CrossSync, "_ReadRowsOperation") as mock_op_constructor:
+                mock_op = mock.Mock()
                 expected_result = object()
-                read_rows.side_effect = lambda *args, **kwargs: [expected_result]
+
+                if CrossSync.is_async:
+
+                    async def mock_generator():
+                        yield expected_result
+
+                    mock_op.start_operation.return_value = mock_generator()
+                else:
+                    mock_op.start_operation.return_value = [expected_result]
+                mock_op_constructor.return_value = mock_op
                 expected_op_timeout = 8
                 expected_req_timeout = 4
                 row = await table.read_row(
@@ -1970,43 +1978,52 @@ async def test_read_row(self):
                     attempt_timeout=expected_req_timeout,
                 )
                 assert row == expected_result
-                assert read_rows.call_count == 1
-                args, kwargs = read_rows.call_args_list[0]
+                assert mock_op_constructor.call_count == 1
+                args, kwargs = mock_op_constructor.call_args_list[0]
                 assert kwargs["operation_timeout"] == expected_op_timeout
                 assert kwargs["attempt_timeout"] == expected_req_timeout
-                assert len(args) == 1
+                assert len(args) == 2
                 assert isinstance(args[0], ReadRowsQuery)
                 query = args[0]
                 assert query.row_keys == [row_key]
                 assert query.row_ranges == []
                 assert query.limit == 1
+                assert args[1] is table
 
     @CrossSync.pytest
     async def test_read_row_w_filter(self):
         """Test reading a single row with an added filter"""
         async with self._make_client() as client:
             table = client.get_table("instance", "table")
             row_key = b"test_1"
-            with mock.patch.object(table, "read_rows") as read_rows:
+            with mock.patch.object(CrossSync, "_ReadRowsOperation") as mock_op_constructor:
+                mock_op = mock.Mock()
                 expected_result = object()
-                read_rows.side_effect = lambda *args, **kwargs: [expected_result]
+
+                if CrossSync.is_async:
+
+                    async def mock_generator():
+                        yield expected_result
+
+                    mock_op.start_operation.return_value = mock_generator()
+                else:
+                    mock_op.start_operation.return_value = [expected_result]
+                mock_op_constructor.return_value = mock_op
                 expected_op_timeout = 8
                 expected_req_timeout = 4
-                mock_filter = mock.Mock()
-                expected_filter = {"filter": "mock filter"}
-                mock_filter._to_dict.return_value = expected_filter
+                expected_filter = mock.Mock()
                 row = await table.read_row(
                     row_key,
                     operation_timeout=expected_op_timeout,
                     attempt_timeout=expected_req_timeout,
                     row_filter=expected_filter,
                 )
                 assert row == expected_result
-                assert read_rows.call_count == 1
-                args, kwargs = read_rows.call_args_list[0]
+                assert mock_op_constructor.call_count == 1
+                args, kwargs = mock_op_constructor.call_args_list[0]
                 assert kwargs["operation_timeout"] == expected_op_timeout
                 assert kwargs["attempt_timeout"] == expected_req_timeout
-                assert len(args) == 1
+                assert len(args) == 2
                 assert isinstance(args[0], ReadRowsQuery)
                 query = args[0]
                 assert query.row_keys == [row_key]
@@ -2020,9 +2037,19 @@ async def test_read_row_no_response(self):
         async with self._make_client() as client:
             table = client.get_table("instance", "table")
             row_key = b"test_1"
-            with mock.patch.object(table, "read_rows") as read_rows:
-                # return no rows
-                read_rows.side_effect = lambda *args, **kwargs: []
+            with mock.patch.object(CrossSync, "_ReadRowsOperation") as mock_op_constructor:
+                mock_op = mock.Mock()
+
+                if CrossSync.is_async:
+
+                    async def mock_generator():
+                        if False:
+                            yield
+
+                    mock_op.start_operation.return_value = mock_generator()
+                else:
+                    mock_op.start_operation.return_value = []
+                mock_op_constructor.return_value = mock_op
                 expected_op_timeout = 8
                 expected_req_timeout = 4
                 result = await table.read_row(
@@ -2031,8 +2058,8 @@ async def test_read_row_no_response(self):
                     attempt_timeout=expected_req_timeout,
                 )
                 assert result is None
-                assert read_rows.call_count == 1
-                args, kwargs = read_rows.call_args_list[0]
+                assert mock_op_constructor.call_count == 1
+                args, kwargs = mock_op_constructor.call_args_list[0]
                 assert kwargs["operation_timeout"] == expected_op_timeout
                 assert kwargs["attempt_timeout"] == expected_req_timeout
                 assert isinstance(args[0], ReadRowsQuery)
@@ -2055,22 +2082,34 @@ async def test_row_exists(self, return_value, expected_result):
         async with self._make_client() as client:
             table = client.get_table("instance", "table")
             row_key = b"test_1"
-            with mock.patch.object(table, "read_rows") as read_rows:
-                # return no rows
-                read_rows.side_effect = lambda *args, **kwargs: return_value
-                expected_op_timeout = 1
-                expected_req_timeout = 2
+            with mock.patch.object(CrossSync, "_ReadRowsOperation") as mock_op_constructor:
+                mock_op = mock.Mock()
+                if CrossSync.is_async:
+
+                    async def mock_generator():
+                        for item in return_value:
+                            yield item
+
+                    mock_op.start_operation.return_value = mock_generator()
+                else:
+                    mock_op.start_operation.return_value = return_value
+                mock_op_constructor.return_value = mock_op
+                expected_op_timeout = 2
+                expected_req_timeout = 1
                 result = await table.row_exists(
                     row_key,
                     operation_timeout=expected_op_timeout,
                     attempt_timeout=expected_req_timeout,
                 )
                 assert expected_result == result
-                assert read_rows.call_count == 1
-                args, kwargs = read_rows.call_args_list[0]
+                assert mock_op_constructor.call_count == 1
+                args, kwargs = mock_op_constructor.call_args_list[0]
                 assert kwargs["operation_timeout"] == expected_op_timeout
                 assert kwargs["attempt_timeout"] == expected_req_timeout
-                assert isinstance(args[0], ReadRowsQuery)
+                query = args[0]
+                assert isinstance(query, ReadRowsQuery)
+                assert query.row_keys == [row_key]
+                assert query.limit == 1
                 expected_filter = {
                     "chain": {
                         "filters": [
@@ -2079,10 +2118,6 @@ async def test_row_exists(self, return_value, expected_result):
                         ]
                     }
                 }
-                query = args[0]
-                assert query.row_keys == [row_key]
-                assert query.row_ranges == []
-                assert query.limit == 1
                 assert query.filter._to_dict() == expected_filter