added _mutate_rows tests

tests/unit/test__mutate_rows.py

@@ -0,0 +1,148 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import pytest
+import sys
+
+import google.cloud.bigtable.exceptions as bigtable_exceptions
+from google.cloud.bigtable_v2.types import MutateRowsResponse
+from google.rpc import status_pb2
+
+# try/except added for compatibility with python < 3.8
+try:
+    from unittest import mock
+    from unittest.mock import AsyncMock  # type: ignore
+except ImportError:  # pragma: NO COVER
+    import mock  # type: ignore
+    from mock import AsyncMock  # type: ignore
+
+class Test_MutateRowsRetryableAttempt():
+
+    async def _mock_stream(self, mutation_dict, error_dict):
+        for idx, entry in mutation_dict.items():
+            code = error_dict.get(idx, 0)
+            yield MutateRowsResponse(
+                entries=[
+                    MutateRowsResponse.Entry(
+                        index=idx,
+                        status=status_pb2.Status(code=code)
+                    )
+                ]
+            )
+
+    def _make_mock_client(self, mutation_dict, error_dict=None):
+        client = mock.Mock()
+        client.mutate_rows = AsyncMock()
+        if error_dict is None:
+            error_dict = {}
+        client.mutate_rows.side_effect = lambda *args, **kwargs: self._mock_stream(mutation_dict, error_dict)
+        return client
+
+    @pytest.mark.asyncio
+    async def test_single_entry_success(self):
+        """Test mutating a single entry"""
+        from google.cloud.bigtable._mutate_rows import _mutate_rows_retryable_attempt
+        mutation = mock.Mock()
+        mutations = {0: mutation}
+        client = self._make_mock_client(mutations)
+        errors = {0: []}
+        expected_request = {"test": "data"}
+        expected_timeout = 9
+        await _mutate_rows_retryable_attempt(client, expected_request, expected_timeout, mutations, errors, lambda x: False)
+        assert mutations[0] is None
+        assert errors[0] == []
+        assert client.mutate_rows.call_count == 1
+        args, kwargs = client.mutate_rows.call_args
+        assert kwargs['timeout'] == expected_timeout
+        assert args[0]["test"] == "data"
+        assert args[0]["entries"] == [mutation._to_dict()]
+
+    @pytest.mark.asyncio
+    async def test_empty_request(self):
+        """Calling with no mutations should result in no calls"""
+        from google.cloud.bigtable._mutate_rows import _mutate_rows_retryable_attempt
+        client = self._make_mock_client({})
+        await _mutate_rows_retryable_attempt(client, {}, None, {}, {}, lambda x: False)
+        assert client.mutate_rows.call_count == 0
+
+    @pytest.mark.asyncio
+    async def test_eventual_success(self):
+        """Fail at first, but eventually succeed"""
+        from google.cloud.bigtable._mutate_rows import _mutate_rows_retryable_attempt
+        mutation = mock.Mock()
+        mutations = {0: mutation}
+        client = self._make_mock_client(mutations)
+        # fail 10 times, then succeed
+        client.mutate_rows.side_effect = [self._mock_stream(mutations, {0: 400})] * 10 + [self._mock_stream(mutations, {})]
+        errors = {0: []}
+        expected_request = {}
+        expected_timeout = 9
+        await _mutate_rows_retryable_attempt(client, expected_request, expected_timeout, mutations, errors, lambda x: True)
+        assert mutations[0] is None
+        assert errors[0] == []
+        assert client.mutate_rows.call_count == 11
+        args, kwargs = client.mutate_rows.call_args
+        assert kwargs['timeout'] == expected_timeout
+        assert args[0]["entries"] == [mutation._to_dict()]
+
+    @pytest.mark.asyncio
+    async def test_partial_success(self):
+        """Some entries succeed, but one fails. Should report the proper index"""
+        from google.cloud.bigtable._mutate_rows import _mutate_rows_retryable_attempt
+        success_mutation = mock.Mock()
+        success_mutation_2 = mock.Mock()
+        failure_mutation = mock.Mock()
+        mutations = {0: success_mutation, 1: failure_mutation, 2: success_mutation_2}
+        errors = {0: [], 1: [], 2: []}
+        client = self._make_mock_client(mutations, error_dict={1: 400})
+        # raise retryable error 3 times, then raise non-retryable error
+        client.mutate_rows.side_effect =  [self._mock_stream(mutations, {1: 300}), self._mock_stream({0: failure_mutation}, {0: 400}), self._mock_stream({0: failure_mutation}, {0: 400}), self._mock_stream({0: failure_mutation}, {0: 500})] * 2
+        expected_request = {}
+        expected_timeout = 9
+        await _mutate_rows_retryable_attempt(client, expected_request, expected_timeout, mutations, errors, lambda x: x.grpc_status_code <= 400)
+        assert mutations == {0: None, 1: None, 2: None}
+        assert errors[0] == []
+        assert errors[2] == []
+        # assert len(errors[1]) == 4
+        assert errors[1][0].grpc_status_code == 300
+        assert errors[1][1].grpc_status_code == 400
+        assert errors[1][2].grpc_status_code == 400
+        assert errors[1][3].grpc_status_code == 500
+
+
+    @pytest.mark.asyncio
+    async def test_retryable_entry_error(self):
+        """Should continue retrying if the error is retryable"""
+        from google.cloud.bigtable._mutate_rows import _mutate_rows_retryable_attempt
+        import asyncio
+        mutation = mock.Mock()
+        mutations = {0: mutation}
+        errors = {0: []}
+        client = self._make_mock_client(mutations, error_dict={0: 400})
+        # raise retryable error 3 times, then raise non-retryable error
+        expected_request = {}
+        expected_timeout = 9
+        with mock.patch.object(asyncio, 'sleep') as mock_sleep:
+            fail_after = 100
+            mock_sleep.side_effect = [None] * fail_after + [asyncio.CancelledError("Cancelled")]
+            try:
+                await _mutate_rows_retryable_attempt(client, expected_request, expected_timeout, mutations, errors, lambda x: True)
+            except asyncio.CancelledError:
+                pass
+            assert mock_sleep.call_count == fail_after + 1
+            assert client.mutate_rows.call_count == fail_after
+            assert len(errors[0]) == fail_after
+
+