fixed merge loss

googleapis · Nov 22, 2024 · 6e38d43 · 6e38d43
1 parent a8b93cf
commit 6e38d43
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py
@@ -18,6 +18,7 @@
 
 from google.api_core import exceptions as core_exceptions
 from google.api_core import retry as retries
+import google.cloud.bigtable_v2.types.bigtable as types_pb
 import google.cloud.bigtable.data.exceptions as bt_exceptions
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
 from google.cloud.bigtable.data._helpers import _retry_exception_factory

diff --git a/tests/unit/data/_async/test__mutate_rows.py b/tests/unit/data/_async/test__mutate_rows.py
@@ -52,9 +52,8 @@ def _make_one(self, *args, **kwargs):
         return self._target_class()(*args, **kwargs)
 
     def _make_mutation(self, count=1, size=1):
-        mutation = mock.Mock()
-        mutation.size.return_value = size
-        mutation.mutations = [mock.Mock()] * count
+        mutation = RowMutationEntry("k", [DeleteAllFromRow() for _ in range(count)])
+        mutation.size = lambda: size
         return mutation
 
     @CrossSync.convert

diff --git a/tests/unit/data/_async/test_mutations_batcher.py b/tests/unit/data/_async/test_mutations_batcher.py
@@ -40,9 +40,9 @@ def _make_one(self, max_mutation_count=10, max_mutation_bytes=100):
 
     @staticmethod
     def _make_mutation(count=1, size=1):
-        mutation = mock.Mock()
-        mutation.size.return_value = size
-        mutation.mutations = [mock.Mock()] * count
+        mutation = RowMutationEntry("k", DeleteAllFromRow())
+        mutation.mutations = [DeleteAllFromRow() for _ in range(count)]
+        mutation.size = lambda: size
         return mutation
 
     def test_ctor(self):