test_client.py

# Copyright 2024 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.

# This file is automatically generated by CrossSync. Do not edit manually.

from __future__ import annotations
import grpc
import asyncio
import re
import pytest
import mock
from google.cloud.bigtable.data import mutations
from google.auth.credentials import AnonymousCredentials
from google.cloud.bigtable_v2.types import ReadRowsResponse
from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
from google.api_core import exceptions as core_exceptions
from google.api_core import client_options
from google.cloud.bigtable.data.exceptions import InvalidChunk
from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete
from google.cloud.bigtable.data.mutations import DeleteAllFromRow
from google.cloud.bigtable.data import TABLE_DEFAULT
from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule
from google.cloud.bigtable.data.read_modify_write_rules import AppendValueRule
from google.cloud.bigtable_v2.types.bigtable import ExecuteQueryResponse
from google.cloud.bigtable.data._cross_sync import CrossSync
from tests.unit.data.execute_query.sql_helpers import (
    chunked_responses,
    column,
    int64_type,
    int_val,
    metadata,
    null_val,
    prepare_response,
    str_type,
    str_val,
)
from google.api_core import grpc_helpers

CrossSync._Sync_Impl.add_mapping("grpc_helpers", grpc_helpers)


@CrossSync._Sync_Impl.add_mapping_decorator("TestBigtableDataClient")
class TestBigtableDataClient:
    @staticmethod
    def _get_target_class():
        return CrossSync._Sync_Impl.DataClient

    @classmethod
    def _make_client(cls, *args, use_emulator=True, **kwargs):
        import os

        env_mask = {}
        if use_emulator:
            env_mask["BIGTABLE_EMULATOR_HOST"] = "localhost"
            import warnings

            warnings.filterwarnings("ignore", category=RuntimeWarning)
        else:
            kwargs["credentials"] = kwargs.get("credentials", AnonymousCredentials())
            kwargs["project"] = kwargs.get("project", "project-id")
        with mock.patch.dict(os.environ, env_mask):
            return cls._get_target_class()(*args, **kwargs)

    def test_ctor(self):
        expected_project = "project-id"
        expected_credentials = AnonymousCredentials()
        client = self._make_client(
            project="project-id", credentials=expected_credentials, use_emulator=False
        )
        CrossSync._Sync_Impl.yield_to_event_loop()
        assert client.project == expected_project
        assert not client._active_instances
        assert client._channel_refresh_task is not None
        assert client.transport._credentials == expected_credentials
        client.close()

    def test_ctor_super_inits(self):
        from google.cloud.client import ClientWithProject
        from google.api_core import client_options as client_options_lib

        project = "project-id"
        credentials = AnonymousCredentials()
        client_options = {"api_endpoint": "foo.bar:1234"}
        options_parsed = client_options_lib.from_dict(client_options)
        with mock.patch.object(
            CrossSync._Sync_Impl.GapicClient, "__init__"
        ) as bigtable_client_init:
            bigtable_client_init.return_value = None
            with mock.patch.object(
                ClientWithProject, "__init__"
            ) as client_project_init:
                client_project_init.return_value = None
                try:
                    self._make_client(
                        project=project,
                        credentials=credentials,
                        client_options=options_parsed,
                        use_emulator=False,
                    )
                except AttributeError:
                    pass
                assert bigtable_client_init.call_count == 1
                kwargs = bigtable_client_init.call_args[1]
                assert kwargs["credentials"] == credentials
                assert kwargs["client_options"] == options_parsed
                assert client_project_init.call_count == 1
                kwargs = client_project_init.call_args[1]
                assert kwargs["project"] == project
                assert kwargs["credentials"] == credentials
                assert kwargs["client_options"] == options_parsed

    def test_ctor_dict_options(self):
        from google.api_core.client_options import ClientOptions

        client_options = {"api_endpoint": "foo.bar:1234"}
        with mock.patch.object(
            CrossSync._Sync_Impl.GapicClient, "__init__"
        ) as bigtable_client_init:
            try:
                self._make_client(client_options=client_options)
            except TypeError:
                pass
            bigtable_client_init.assert_called_once()
            kwargs = bigtable_client_init.call_args[1]
            called_options = kwargs["client_options"]
            assert called_options.api_endpoint == "foo.bar:1234"
            assert isinstance(called_options, ClientOptions)
        with mock.patch.object(
            self._get_target_class(), "_start_background_channel_refresh"
        ) as start_background_refresh:
            client = self._make_client(
                client_options=client_options, use_emulator=False
            )
            start_background_refresh.assert_called_once()
            client.close()

    def test_veneer_grpc_headers(self):
        client_component = "data-async" if CrossSync._Sync_Impl.is_async else "data"
        VENEER_HEADER_REGEX = re.compile(
            "gapic\\/[0-9]+\\.[\\w.-]+ gax\\/[0-9]+\\.[\\w.-]+ gccl\\/[0-9]+\\.[\\w.-]+-"
            + client_component
            + " gl-python\\/[0-9]+\\.[\\w.-]+ grpc\\/[0-9]+\\.[\\w.-]+"
        )
        patch = mock.patch("google.api_core.gapic_v1.method.wrap_method")
        with patch as gapic_mock:
            client = self._make_client(project="project-id")
            wrapped_call_list = gapic_mock.call_args_list
            assert len(wrapped_call_list) > 0
            for call in wrapped_call_list:
                client_info = call.kwargs["client_info"]
                assert client_info is not None, f"{call} has no client_info"
                wrapped_user_agent_sorted = " ".join(
                    sorted(client_info.to_user_agent().split(" "))
                )
                assert VENEER_HEADER_REGEX.match(
                    wrapped_user_agent_sorted
                ), f"'{wrapped_user_agent_sorted}' does not match {VENEER_HEADER_REGEX}"
            client.close()

    def test__start_background_channel_refresh_task_exists(self):
        client = self._make_client(project="project-id", use_emulator=False)
        assert client._channel_refresh_task is not None
        with mock.patch.object(asyncio, "create_task") as create_task:
            client._start_background_channel_refresh()
            create_task.assert_not_called()
        client.close()

    def test__start_background_channel_refresh(self):
        client = self._make_client(project="project-id")
        with mock.patch.object(
            client, "_ping_and_warm_instances", CrossSync._Sync_Impl.Mock()
        ) as ping_and_warm:
            client._emulator_host = None
            client._start_background_channel_refresh()
            assert client._channel_refresh_task is not None
            assert isinstance(client._channel_refresh_task, CrossSync._Sync_Impl.Task)
            CrossSync._Sync_Impl.sleep(0.1)
            assert ping_and_warm.call_count == 1
            client.close()

    def test__ping_and_warm_instances(self):
        """test ping and warm with mocked asyncio.gather"""
        client_mock = mock.Mock()
        client_mock._execute_ping_and_warms = (
            lambda *args: self._get_target_class()._execute_ping_and_warms(
                client_mock, *args
            )
        )
        with mock.patch.object(
            CrossSync._Sync_Impl, "gather_partials", CrossSync._Sync_Impl.Mock()
        ) as gather:
            gather.side_effect = lambda partials, **kwargs: [None for _ in partials]
            channel = mock.Mock()
            client_mock._active_instances = []
            result = self._get_target_class()._ping_and_warm_instances(
                client_mock, channel=channel
            )
            assert len(result) == 0
            assert gather.call_args[1]["return_exceptions"] is True
            assert gather.call_args[1]["sync_executor"] == client_mock._executor
            client_mock._active_instances = [(mock.Mock(), mock.Mock())] * 4
            gather.reset_mock()
            channel.reset_mock()
            result = self._get_target_class()._ping_and_warm_instances(
                client_mock, channel=channel
            )
            assert len(result) == 4
            gather.assert_called_once()
            partial_list = gather.call_args.args[0]
            assert len(partial_list) == 4
            grpc_call_args = channel.unary_unary().call_args_list
            for idx, (_, kwargs) in enumerate(grpc_call_args):
                (
                    expected_instance,
                    expected_app_profile,
                ) = client_mock._active_instances[idx]
                request = kwargs["request"]
                assert request["name"] == expected_instance
                assert request["app_profile_id"] == expected_app_profile
                metadata = kwargs["metadata"]
                assert len(metadata) == 1
                assert metadata[0][0] == "x-goog-request-params"
                assert (
                    metadata[0][1]
                    == f"name={expected_instance}&app_profile_id={expected_app_profile}"
                )

    def test__ping_and_warm_single_instance(self):
        """should be able to call ping and warm with single instance"""
        client_mock = mock.Mock()
        client_mock._execute_ping_and_warms = (
            lambda *args: self._get_target_class()._execute_ping_and_warms(
                client_mock, *args
            )
        )
        with mock.patch.object(
            CrossSync._Sync_Impl, "gather_partials", CrossSync._Sync_Impl.Mock()
        ) as gather:
            gather.side_effect = lambda *args, **kwargs: [fn() for fn in args[0]]
            client_mock._active_instances = [mock.Mock()] * 100
            test_key = ("test-instance", "test-app-profile")
            result = self._get_target_class()._ping_and_warm_instances(
                client_mock, test_key
            )
            assert len(result) == 1
            grpc_call_args = (
                client_mock.transport.grpc_channel.unary_unary().call_args_list
            )
            assert len(grpc_call_args) == 1
            kwargs = grpc_call_args[0][1]
            request = kwargs["request"]
            assert request["name"] == "test-instance"
            assert request["app_profile_id"] == "test-app-profile"
            metadata = kwargs["metadata"]
            assert len(metadata) == 1
            assert metadata[0][0] == "x-goog-request-params"
            assert (
                metadata[0][1] == "name=test-instance&app_profile_id=test-app-profile"
            )

    @pytest.mark.parametrize(
        "refresh_interval, wait_time, expected_sleep",
        [(0, 0, 0), (0, 1, 0), (10, 0, 10), (10, 5, 5), (10, 10, 0), (10, 15, 0)],
    )
    def test__manage_channel_first_sleep(
        self, refresh_interval, wait_time, expected_sleep
    ):
        import time

        with mock.patch.object(time, "monotonic") as monotonic:
            monotonic.return_value = 0
            with mock.patch.object(CrossSync._Sync_Impl, "event_wait") as sleep:
                sleep.side_effect = asyncio.CancelledError
                try:
                    client = self._make_client(project="project-id")
                    client._channel_init_time = -wait_time
                    client._manage_channel(refresh_interval, refresh_interval)
                except asyncio.CancelledError:
                    pass
                sleep.assert_called_once()
                call_time = sleep.call_args[0][1]
                assert (
                    abs(call_time - expected_sleep) < 0.1
                ), f"refresh_interval: {refresh_interval}, wait_time: {wait_time}, expected_sleep: {expected_sleep}"
                client.close()

    def test__manage_channel_ping_and_warm(self):
        """_manage channel should call ping and warm internally"""
        import time
        import threading
        from google.cloud.bigtable_v2.services.bigtable.transports.grpc import (
            _LoggingClientInterceptor as Interceptor,
        )

        client_mock = mock.Mock()
        client_mock.transport._interceptor = Interceptor()
        client_mock._is_closed.is_set.return_value = False
        client_mock._channel_init_time = time.monotonic()
        orig_channel = client_mock.transport.grpc_channel
        sleep_tuple = (
            (asyncio, "sleep")
            if CrossSync._Sync_Impl.is_async
            else (threading.Event, "wait")
        )
        with mock.patch.object(*sleep_tuple):
            orig_channel.close.side_effect = asyncio.CancelledError
            ping_and_warm = (
                client_mock._ping_and_warm_instances
            ) = CrossSync._Sync_Impl.Mock()
            try:
                self._get_target_class()._manage_channel(client_mock, 10)
            except asyncio.CancelledError:
                pass
            assert ping_and_warm.call_count == 2
            assert client_mock.transport._grpc_channel != orig_channel
            called_with = [call[1]["channel"] for call in ping_and_warm.call_args_list]
            assert orig_channel in called_with
            assert client_mock.transport.grpc_channel in called_with

    @pytest.mark.parametrize(
        "refresh_interval, num_cycles, expected_sleep",
        [(None, 1, 60 * 35), (10, 10, 100), (10, 1, 10)],
    )
    def test__manage_channel_sleeps(self, refresh_interval, num_cycles, expected_sleep):
        import time
        import random

        channel = mock.Mock()
        channel.close = CrossSync._Sync_Impl.Mock()
        with mock.patch.object(random, "uniform") as uniform:
            uniform.side_effect = lambda min_, max_: min_
            with mock.patch.object(time, "time") as time_mock:
                time_mock.return_value = 0
                with mock.patch.object(CrossSync._Sync_Impl, "event_wait") as sleep:
                    sleep.side_effect = [None for i in range(num_cycles - 1)] + [
                        asyncio.CancelledError
                    ]
                    client = self._make_client(project="project-id")
                    client.transport._grpc_channel = channel
                    with mock.patch.object(
                        client.transport, "create_channel", CrossSync._Sync_Impl.Mock
                    ):
                        try:
                            if refresh_interval is not None:
                                client._manage_channel(
                                    refresh_interval, refresh_interval, grace_period=0
                                )
                            else:
                                client._manage_channel(grace_period=0)
                        except asyncio.CancelledError:
                            pass
                    assert sleep.call_count == num_cycles
                    total_sleep = sum([call[0][1] for call in sleep.call_args_list])
                    assert (
                        abs(total_sleep - expected_sleep) < 0.5
                    ), f"refresh_interval={refresh_interval}, num_cycles={num_cycles}, expected_sleep={expected_sleep}"
        client.close()

    def test__manage_channel_random(self):
        import random

        with mock.patch.object(CrossSync._Sync_Impl, "event_wait") as sleep:
            with mock.patch.object(random, "uniform") as uniform:
                uniform.return_value = 0
                try:
                    uniform.side_effect = asyncio.CancelledError
                    client = self._make_client(project="project-id")
                except asyncio.CancelledError:
                    uniform.side_effect = None
                    uniform.reset_mock()
                    sleep.reset_mock()
                with mock.patch.object(client.transport, "create_channel"):
                    min_val = 200
                    max_val = 205
                    uniform.side_effect = lambda min_, max_: min_
                    sleep.side_effect = [None, asyncio.CancelledError]
                    try:
                        client._manage_channel(min_val, max_val, grace_period=0)
                    except asyncio.CancelledError:
                        pass
                    assert uniform.call_count == 2
                    uniform_args = [call[0] for call in uniform.call_args_list]
                    for found_min, found_max in uniform_args:
                        assert found_min == min_val
                        assert found_max == max_val

    @pytest.mark.parametrize("num_cycles", [0, 1, 10, 100])
    def test__manage_channel_refresh(self, num_cycles):
        expected_refresh = 0.5
        grpc_lib = grpc.aio if CrossSync._Sync_Impl.is_async else grpc
        new_channel = grpc_lib.insecure_channel("localhost:8080")
        with mock.patch.object(CrossSync._Sync_Impl, "event_wait") as sleep:
            sleep.side_effect = [None for i in range(num_cycles)] + [RuntimeError]
            with mock.patch.object(
                CrossSync._Sync_Impl.grpc_helpers, "create_channel"
            ) as create_channel:
                create_channel.return_value = new_channel
                client = self._make_client(project="project-id")
                create_channel.reset_mock()
                try:
                    client._manage_channel(
                        refresh_interval_min=expected_refresh,
                        refresh_interval_max=expected_refresh,
                        grace_period=0,
                    )
                except RuntimeError:
                    pass
                assert sleep.call_count == num_cycles + 1
                assert create_channel.call_count == num_cycles
            client.close()

    def test__register_instance(self):
        """test instance registration"""
        client_mock = mock.Mock()
        client_mock._gapic_client.instance_path.side_effect = lambda a, b: f"prefix/{b}"
        active_instances = set()
        instance_owners = {}
        client_mock._active_instances = active_instances
        client_mock._instance_owners = instance_owners
        client_mock._channel_refresh_task = None
        client_mock._ping_and_warm_instances = CrossSync._Sync_Impl.Mock()
        table_mock = mock.Mock()
        self._get_target_class()._register_instance(
            client_mock, "instance-1", table_mock
        )
        assert client_mock._start_background_channel_refresh.call_count == 1
        expected_key = ("prefix/instance-1", table_mock.app_profile_id)
        assert len(active_instances) == 1
        assert expected_key == tuple(list(active_instances)[0])
        assert len(instance_owners) == 1
        assert expected_key == tuple(list(instance_owners)[0])
        client_mock._channel_refresh_task = mock.Mock()
        table_mock2 = mock.Mock()
        self._get_target_class()._register_instance(
            client_mock, "instance-2", table_mock2
        )
        assert client_mock._start_background_channel_refresh.call_count == 1
        assert (
            client_mock._ping_and_warm_instances.call_args[0][0][0]
            == "prefix/instance-2"
        )
        assert client_mock._ping_and_warm_instances.call_count == 1
        assert len(active_instances) == 2
        assert len(instance_owners) == 2
        expected_key2 = ("prefix/instance-2", table_mock2.app_profile_id)
        assert any(
            [
                expected_key2 == tuple(list(active_instances)[i])
                for i in range(len(active_instances))
            ]
        )
        assert any(
            [
                expected_key2 == tuple(list(instance_owners)[i])
                for i in range(len(instance_owners))
            ]
        )

    def test__register_instance_duplicate(self):
        """test double instance registration. Should be no-op"""
        client_mock = mock.Mock()
        client_mock._gapic_client.instance_path.side_effect = lambda a, b: f"prefix/{b}"
        active_instances = set()
        instance_owners = {}
        client_mock._active_instances = active_instances
        client_mock._instance_owners = instance_owners
        client_mock._channel_refresh_task = object()
        mock_channels = [mock.Mock()]
        client_mock.transport.channels = mock_channels
        client_mock._ping_and_warm_instances = CrossSync._Sync_Impl.Mock()
        table_mock = mock.Mock()
        expected_key = ("prefix/instance-1", table_mock.app_profile_id)
        self._get_target_class()._register_instance(
            client_mock, "instance-1", table_mock
        )
        assert len(active_instances) == 1
        assert expected_key == tuple(list(active_instances)[0])
        assert len(instance_owners) == 1
        assert expected_key == tuple(list(instance_owners)[0])
        assert client_mock._ping_and_warm_instances.call_count == 1
        self._get_target_class()._register_instance(
            client_mock, "instance-1", table_mock
        )
        assert len(active_instances) == 1
        assert expected_key == tuple(list(active_instances)[0])
        assert len(instance_owners) == 1
        assert expected_key == tuple(list(instance_owners)[0])
        assert client_mock._ping_and_warm_instances.call_count == 1

    @pytest.mark.parametrize(
        "insert_instances,expected_active,expected_owner_keys",
        [
            ([("i", None)], [("i", None)], [("i", None)]),
            ([("i", "p")], [("i", "p")], [("i", "p")]),
            ([("1", "p"), ("1", "p")], [("1", "p")], [("1", "p")]),
            (
                [("1", "p"), ("2", "p")],
                [("1", "p"), ("2", "p")],
                [("1", "p"), ("2", "p")],
            ),
        ],
    )
    def test__register_instance_state(
        self, insert_instances, expected_active, expected_owner_keys
    ):
        """test that active_instances and instance_owners are updated as expected"""
        client_mock = mock.Mock()
        client_mock._gapic_client.instance_path.side_effect = lambda a, b: b
        active_instances = set()
        instance_owners = {}
        client_mock._active_instances = active_instances
        client_mock._instance_owners = instance_owners
        client_mock._channel_refresh_task = None
        client_mock._ping_and_warm_instances = CrossSync._Sync_Impl.Mock()
        table_mock = mock.Mock()
        for instance, profile in insert_instances:
            table_mock.app_profile_id = profile
            self._get_target_class()._register_instance(
                client_mock, instance, table_mock
            )
        assert len(active_instances) == len(expected_active)
        assert len(instance_owners) == len(expected_owner_keys)
        for expected in expected_active:
            assert any(
                [
                    expected == tuple(list(active_instances)[i])
                    for i in range(len(active_instances))
                ]
            )
        for expected in expected_owner_keys:
            assert any(
                [
                    expected == tuple(list(instance_owners)[i])
                    for i in range(len(instance_owners))
                ]
            )

    def test__remove_instance_registration(self):
        client = self._make_client(project="project-id")
        table = mock.Mock()
        client._register_instance("instance-1", table)
        client._register_instance("instance-2", table)
        assert len(client._active_instances) == 2
        assert len(client._instance_owners.keys()) == 2
        instance_1_path = client._gapic_client.instance_path(
            client.project, "instance-1"
        )
        instance_1_key = (instance_1_path, table.app_profile_id)
        instance_2_path = client._gapic_client.instance_path(
            client.project, "instance-2"
        )
        instance_2_key = (instance_2_path, table.app_profile_id)
        assert len(client._instance_owners[instance_1_key]) == 1
        assert list(client._instance_owners[instance_1_key])[0] == id(table)
        assert len(client._instance_owners[instance_2_key]) == 1
        assert list(client._instance_owners[instance_2_key])[0] == id(table)
        success = client._remove_instance_registration("instance-1", table)
        assert success
        assert len(client._active_instances) == 1
        assert len(client._instance_owners[instance_1_key]) == 0
        assert len(client._instance_owners[instance_2_key]) == 1
        assert client._active_instances == {instance_2_key}
        success = client._remove_instance_registration("fake-key", table)
        assert not success
        assert len(client._active_instances) == 1
        client.close()

    def test__multiple_table_registration(self):
        """registering with multiple tables with the same key should
        add multiple owners to instance_owners, but only keep one copy
        of shared key in active_instances"""
        from google.cloud.bigtable.data._helpers import _WarmedInstanceKey

        with self._make_client(project="project-id") as client:
            with client.get_table("instance_1", "table_1") as table_1:
                instance_1_path = client._gapic_client.instance_path(
                    client.project, "instance_1"
                )
                instance_1_key = _WarmedInstanceKey(
                    instance_1_path, table_1.app_profile_id
                )
                assert len(client._instance_owners[instance_1_key]) == 1
                assert len(client._active_instances) == 1
                assert id(table_1) in client._instance_owners[instance_1_key]
                with client.get_table("instance_1", "table_2") as table_2:
                    assert table_2._register_instance_future is not None
                    table_2._register_instance_future.result()
                    assert len(client._instance_owners[instance_1_key]) == 2
                    assert len(client._active_instances) == 1
                    assert id(table_1) in client._instance_owners[instance_1_key]
                    assert id(table_2) in client._instance_owners[instance_1_key]
                    with client.get_table(
                        "instance_1", "table_3", app_profile_id="diff"
                    ) as table_3:
                        assert table_3._register_instance_future is not None
                        table_3._register_instance_future.result()
                        instance_3_path = client._gapic_client.instance_path(
                            client.project, "instance_1"
                        )
                        instance_3_key = _WarmedInstanceKey(
                            instance_3_path, table_3.app_profile_id
                        )
                        assert len(client._instance_owners[instance_1_key]) == 2
                        assert len(client._instance_owners[instance_3_key]) == 1
                        assert len(client._active_instances) == 2
                        assert id(table_1) in client._instance_owners[instance_1_key]
                        assert id(table_2) in client._instance_owners[instance_1_key]
                        assert id(table_3) in client._instance_owners[instance_3_key]
                assert len(client._active_instances) == 1
                assert instance_1_key in client._active_instances
                assert id(table_2) not in client._instance_owners[instance_1_key]
            assert len(client._active_instances) == 0
            assert instance_1_key not in client._active_instances
            assert len(client._instance_owners[instance_1_key]) == 0

    def test__multiple_instance_registration(self):
        """registering with multiple instance keys should update the key
        in instance_owners and active_instances"""
        from google.cloud.bigtable.data._helpers import _WarmedInstanceKey

        with self._make_client(project="project-id") as client:
            with client.get_table("instance_1", "table_1") as table_1:
                assert table_1._register_instance_future is not None
                table_1._register_instance_future.result()
                with client.get_table("instance_2", "table_2") as table_2:
                    assert table_2._register_instance_future is not None
                    table_2._register_instance_future.result()
                    instance_1_path = client._gapic_client.instance_path(
                        client.project, "instance_1"
                    )
                    instance_1_key = _WarmedInstanceKey(
                        instance_1_path, table_1.app_profile_id
                    )
                    instance_2_path = client._gapic_client.instance_path(
                        client.project, "instance_2"
                    )
                    instance_2_key = _WarmedInstanceKey(
                        instance_2_path, table_2.app_profile_id
                    )
                    assert len(client._instance_owners[instance_1_key]) == 1
                    assert len(client._instance_owners[instance_2_key]) == 1
                    assert len(client._active_instances) == 2
                    assert id(table_1) in client._instance_owners[instance_1_key]
                    assert id(table_2) in client._instance_owners[instance_2_key]
                assert len(client._active_instances) == 1
                assert instance_1_key in client._active_instances
                assert len(client._instance_owners[instance_2_key]) == 0
                assert len(client._instance_owners[instance_1_key]) == 1
                assert id(table_1) in client._instance_owners[instance_1_key]
            assert len(client._active_instances) == 0
            assert len(client._instance_owners[instance_1_key]) == 0
            assert len(client._instance_owners[instance_2_key]) == 0

    @pytest.mark.parametrize("method", ["get_table", "get_authorized_view"])
    def test_get_api_surface(self, method):
        """test client.get_table and client.get_authorized_view"""
        from google.cloud.bigtable.data._helpers import _WarmedInstanceKey

        client = self._make_client(project="project-id")
        assert not client._active_instances
        expected_table_id = "table-id"
        expected_instance_id = "instance-id"
        expected_app_profile_id = "app-profile-id"
        if method == "get_table":
            surface = client.get_table(
                expected_instance_id, expected_table_id, expected_app_profile_id
            )
            assert isinstance(
                surface, CrossSync._Sync_Impl.TestTable._get_target_class()
            )
        elif method == "get_authorized_view":
            surface = client.get_authorized_view(
                expected_instance_id,
                expected_table_id,
                "view_id",
                expected_app_profile_id,
            )
            assert isinstance(
                surface, CrossSync._Sync_Impl.TestAuthorizedView._get_target_class()
            )
            assert (
                surface.authorized_view_name
                == f"projects/{client.project}/instances/{expected_instance_id}/tables/{expected_table_id}/authorizedViews/view_id"
            )
        else:
            raise TypeError(f"unexpected method: {method}")
        CrossSync._Sync_Impl.yield_to_event_loop()
        assert surface.table_id == expected_table_id
        assert (
            surface.table_name
            == f"projects/{client.project}/instances/{expected_instance_id}/tables/{expected_table_id}"
        )
        assert surface.instance_id == expected_instance_id
        assert (
            surface.instance_name
            == f"projects/{client.project}/instances/{expected_instance_id}"
        )
        assert surface.app_profile_id == expected_app_profile_id
        assert surface.client is client
        instance_key = _WarmedInstanceKey(surface.instance_name, surface.app_profile_id)
        assert instance_key in client._active_instances
        assert client._instance_owners[instance_key] == {id(surface)}
        client.close()

    @pytest.mark.parametrize("method", ["get_table", "get_authorized_view"])
    def test_api_surface_arg_passthrough(self, method):
        """All arguments passed in get_table and get_authorized_view should be sent to constructor"""
        if method == "get_table":
            surface_type = CrossSync._Sync_Impl.TestTable._get_target_class()
        elif method == "get_authorized_view":
            surface_type = CrossSync._Sync_Impl.TestAuthorizedView._get_target_class()
        else:
            raise TypeError(f"unexpected method: {method}")
        with self._make_client(project="project-id") as client:
            with mock.patch.object(surface_type, "__init__") as mock_constructor:
                mock_constructor.return_value = None
                assert not client._active_instances
                expected_args = (
                    "table",
                    "instance",
                    "view",
                    "app_profile",
                    1,
                    "test",
                    {"test": 2},
                )
                expected_kwargs = {"hello": "world", "test": 2}
                getattr(client, method)(*expected_args, **expected_kwargs)
                mock_constructor.assert_called_once_with(
                    client, *expected_args, **expected_kwargs
                )

    @pytest.mark.parametrize("method", ["get_table", "get_authorized_view"])
    def test_api_surface_context_manager(self, method):
        """get_table and get_authorized_view should work as context managers"""
        from functools import partial
        from google.cloud.bigtable.data._helpers import _WarmedInstanceKey

        expected_table_id = "table-id"
        expected_instance_id = "instance-id"
        expected_app_profile_id = "app-profile-id"
        expected_project_id = "project-id"
        if method == "get_table":
            surface_type = CrossSync._Sync_Impl.TestTable._get_target_class()
        elif method == "get_authorized_view":
            surface_type = CrossSync._Sync_Impl.TestAuthorizedView._get_target_class()
        else:
            raise TypeError(f"unexpected method: {method}")
        with mock.patch.object(surface_type, "close") as close_mock:
            with self._make_client(project=expected_project_id) as client:
                if method == "get_table":
                    fn = partial(
                        client.get_table,
                        expected_instance_id,
                        expected_table_id,
                        expected_app_profile_id,
                    )
                elif method == "get_authorized_view":
                    fn = partial(
                        client.get_authorized_view,
                        expected_instance_id,
                        expected_table_id,
                        "view_id",
                        expected_app_profile_id,
                    )
                else:
                    raise TypeError(f"unexpected method: {method}")
                with fn() as table:
                    CrossSync._Sync_Impl.yield_to_event_loop()
                    assert isinstance(table, surface_type)
                    assert table.table_id == expected_table_id
                    assert (
                        table.table_name
                        == f"projects/{expected_project_id}/instances/{expected_instance_id}/tables/{expected_table_id}"
                    )
                    assert table.instance_id == expected_instance_id
                    assert (
                        table.instance_name
                        == f"projects/{expected_project_id}/instances/{expected_instance_id}"
                    )
                    assert table.app_profile_id == expected_app_profile_id
                    assert table.client is client
                    instance_key = _WarmedInstanceKey(
                        table.instance_name, table.app_profile_id
                    )
                    assert instance_key in client._active_instances
                    assert client._instance_owners[instance_key] == {id(table)}
            assert close_mock.call_count == 1

    def test_close(self):
        client = self._make_client(project="project-id", use_emulator=False)
        task = client._channel_refresh_task
        assert task is not None
        assert not task.done()
        with mock.patch.object(
            client.transport, "close", CrossSync._Sync_Impl.Mock()
        ) as close_mock:
            client.close()
            close_mock.assert_called_once()
        assert task.done()
        assert client._channel_refresh_task is None

    def test_close_with_timeout(self):
        expected_timeout = 19
        client = self._make_client(project="project-id", use_emulator=False)
        with mock.patch.object(
            CrossSync._Sync_Impl, "wait", CrossSync._Sync_Impl.Mock()
        ) as wait_for_mock:
            client.close(timeout=expected_timeout)
            wait_for_mock.assert_called_once()
            assert wait_for_mock.call_args[1]["timeout"] == expected_timeout
        client.close()

    def test_context_manager(self):
        from functools import partial

        close_mock = CrossSync._Sync_Impl.Mock()
        true_close = None
        with self._make_client(project="project-id", use_emulator=False) as client:
            true_close = partial(client.close)
            client.close = close_mock
            assert not client._channel_refresh_task.done()
            assert client.project == "project-id"
            assert client._active_instances == set()
            close_mock.assert_not_called()
        close_mock.assert_called_once()
        true_close()

    def test_default_universe_domain(self):
        """When not passed, universe_domain should default to googleapis.com"""
        with self._make_client(project="project-id", credentials=None) as client:
            assert client.universe_domain == "googleapis.com"
            assert client.api_endpoint == "bigtable.googleapis.com"

    def test_custom_universe_domain(self):
        """test with a customized universe domain value and emulator enabled"""
        universe_domain = "test-universe.test"
        options = client_options.ClientOptions(universe_domain=universe_domain)
        with self._make_client(
            project="project_id",
            client_options=options,
            use_emulator=True,
            credentials=None,
        ) as client:
            assert client.universe_domain == universe_domain
            assert client.api_endpoint == f"bigtable.{universe_domain}"

    def test_configured_universe_domain_matches_GDU(self):
        """that configured universe domain succeeds with matched GDU credentials."""
        universe_domain = "googleapis.com"
        options = client_options.ClientOptions(universe_domain=universe_domain)
        with self._make_client(
            project="project_id", client_options=options, credentials=None
        ) as client:
            assert client.universe_domain == "googleapis.com"
            assert client.api_endpoint == "bigtable.googleapis.com"

    def test_credential_universe_domain_matches_GDU(self):
        """Test with credentials"""
        creds = AnonymousCredentials()
        creds._universe_domain = "googleapis.com"
        with self._make_client(project="project_id", credentials=creds) as client:
            assert client.universe_domain == "googleapis.com"
            assert client.api_endpoint == "bigtable.googleapis.com"

    def test_anomynous_credential_universe_domain(self):
        """Anomynopus credentials should use default universe domain"""
        creds = AnonymousCredentials()
        with self._make_client(project="project_id", credentials=creds) as client:
            assert client.universe_domain == "googleapis.com"
            assert client.api_endpoint == "bigtable.googleapis.com"

    def test_configured_universe_domain_mismatched_credentials(self):
        """Test that configured universe domain errors with mismatched universe
        domain credentials."""
        universe_domain = "test-universe.test"
        options = client_options.ClientOptions(universe_domain=universe_domain)
        creds = AnonymousCredentials()
        creds._universe_domain = "different-universe"
        with pytest.raises(ValueError) as exc:
            self._make_client(
                project="project_id",
                client_options=options,
                use_emulator=False,
                credentials=creds,
            )
        err_msg = f"The configured universe domain ({universe_domain}) does not match the universe domain found in the credentials ({creds.universe_domain}). If you haven't configured the universe domain explicitly, `googleapis.com` is the default."
        assert exc.value.args[0] == err_msg

    def test_configured_universe_domain_matches_credentials(self):
        """Test that configured universe domain succeeds with matching universe
        domain credentials."""
        universe_domain = "test-universe.test"
        options = client_options.ClientOptions(universe_domain=universe_domain)
        creds = AnonymousCredentials()
        creds._universe_domain = universe_domain
        with self._make_client(
            project="project_id", credentials=creds, client_options=options
        ) as client:
            assert client.universe_domain == universe_domain
            assert client.api_endpoint == f"bigtable.{universe_domain}"


@CrossSync._Sync_Impl.add_mapping_decorator("TestTable")
class TestTable:
    def _make_client(self, *args, **kwargs):
        return CrossSync._Sync_Impl.TestBigtableDataClient._make_client(*args, **kwargs)

    @staticmethod
    def _get_target_class():
        return CrossSync._Sync_Impl.Table

    def _make_one(
        self,
        client,
        instance_id="instance",
        table_id="table",
        app_profile_id=None,
        **kwargs,
    ):
        return self._get_target_class()(
            client, instance_id, table_id, app_profile_id, **kwargs
        )

    def test_ctor(self):
        from google.cloud.bigtable.data._helpers import _WarmedInstanceKey

        expected_table_id = "table-id"
        expected_instance_id = "instance-id"
        expected_app_profile_id = "app-profile-id"
        expected_operation_timeout = 123
        expected_attempt_timeout = 12
        expected_read_rows_operation_timeout = 1.5
        expected_read_rows_attempt_timeout = 0.5
        expected_mutate_rows_operation_timeout = 2.5
        expected_mutate_rows_attempt_timeout = 0.75
        client = self._make_client()
        assert not client._active_instances
        table = self._get_target_class()(
            client,
            expected_instance_id,
            expected_table_id,
            expected_app_profile_id,
            default_operation_timeout=expected_operation_timeout,
            default_attempt_timeout=expected_attempt_timeout,
            default_read_rows_operation_timeout=expected_read_rows_operation_timeout,
            default_read_rows_attempt_timeout=expected_read_rows_attempt_timeout,
            default_mutate_rows_operation_timeout=expected_mutate_rows_operation_timeout,
            default_mutate_rows_attempt_timeout=expected_mutate_rows_attempt_timeout,
        )
        CrossSync._Sync_Impl.yield_to_event_loop()
        assert table.table_id == expected_table_id
        assert table.instance_id == expected_instance_id
        assert (
            table.table_name
            == f"projects/{client.project}/instances/{expected_instance_id}/tables/{expected_table_id}"
        )
        assert (
            table.instance_name
            == f"projects/{client.project}/instances/{expected_instance_id}"
        )
        assert table.app_profile_id == expected_app_profile_id
        assert table.client is client
        instance_key = _WarmedInstanceKey(table.instance_name, table.app_profile_id)
        assert instance_key in client._active_instances
        assert client._instance_owners[instance_key] == {id(table)}
        assert table.default_operation_timeout == expected_operation_timeout
        assert table.default_attempt_timeout == expected_attempt_timeout
        assert (
            table.default_read_rows_operation_timeout
            == expected_read_rows_operation_timeout
        )
        assert (
            table.default_read_rows_attempt_timeout
            == expected_read_rows_attempt_timeout
        )
        assert (
            table.default_mutate_rows_operation_timeout
            == expected_mutate_rows_operation_timeout
        )
        assert (
            table.default_mutate_rows_attempt_timeout
            == expected_mutate_rows_attempt_timeout
        )
        table._register_instance_future
        assert table._register_instance_future.done()
        assert not table._register_instance_future.cancelled()
        assert table._register_instance_future.exception() is None
        client.close()

    def test_ctor_defaults(self):
        """should provide default timeout values and app_profile_id"""
        client = self._make_client()
        assert not client._active_instances
        table = self._make_one(client)
        CrossSync._Sync_Impl.yield_to_event_loop()
        assert table.app_profile_id is None
        assert table.client is client
        assert table.default_operation_timeout == 60
        assert table.default_read_rows_operation_timeout == 600
        assert table.default_mutate_rows_operation_timeout == 600
        assert table.default_attempt_timeout == 20
        assert table.default_read_rows_attempt_timeout == 20
        assert table.default_mutate_rows_attempt_timeout == 60
        client.close()

    def test_ctor_invalid_timeout_values(self):
        """bad timeout values should raise ValueError"""
        client = self._make_client()
        timeout_pairs = [
            ("default_operation_timeout", "default_attempt_timeout"),
            (
                "default_read_rows_operation_timeout",
                "default_read_rows_attempt_timeout",
            ),
            (
                "default_mutate_rows_operation_timeout",
                "default_mutate_rows_attempt_timeout",
            ),
        ]
        for operation_timeout, attempt_timeout in timeout_pairs:
            with pytest.raises(ValueError) as e:
                self._make_one(client, **{attempt_timeout: -1})
            assert "attempt_timeout must be greater than 0" in str(e.value)
            with pytest.raises(ValueError) as e:
                self._make_one(client, **{operation_timeout: -1})
            assert "operation_timeout must be greater than 0" in str(e.value)
        client.close()

    @pytest.mark.parametrize(
        "fn_name,fn_args,is_stream,extra_retryables",
        [
            ("read_rows_stream", (ReadRowsQuery(),), True, ()),
            ("read_rows", (ReadRowsQuery(),), True, ()),
            ("read_row", (b"row_key",), True, ()),
            ("read_rows_sharded", ([ReadRowsQuery()],), True, ()),
            ("row_exists", (b"row_key",), True, ()),
            ("sample_row_keys", (), False, ()),
            ("mutate_row", (b"row_key", [DeleteAllFromRow()]), False, ()),
            (
                "bulk_mutate_rows",
                ([mutations.RowMutationEntry(b"key", [DeleteAllFromRow()])],),
                False,
                (_MutateRowsIncomplete,),
            ),
        ],
    )
    @pytest.mark.parametrize(
        "input_retryables,expected_retryables",
        [
            (
                TABLE_DEFAULT.READ_ROWS,
                [
                    core_exceptions.DeadlineExceeded,
                    core_exceptions.ServiceUnavailable,
                    core_exceptions.Aborted,
                ],
            ),
            (
                TABLE_DEFAULT.DEFAULT,
                [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable],
            ),
            (
                TABLE_DEFAULT.MUTATE_ROWS,
                [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable],
            ),
            ([], []),
            ([4], [core_exceptions.DeadlineExceeded]),
        ],
    )
    def test_customizable_retryable_errors(
        self,
        input_retryables,
        expected_retryables,
        fn_name,
        fn_args,
        is_stream,
        extra_retryables,
    ):
        """Test that retryable functions support user-configurable arguments, and that the configured retryables are passed
        down to the gapic layer."""
        retry_fn = "retry_target"
        if is_stream:
            retry_fn += "_stream"
        retry_fn = f"CrossSync._Sync_Impl.{retry_fn}"
        with mock.patch(
            f"google.cloud.bigtable.data._cross_sync.{retry_fn}"
        ) as retry_fn_mock:
            with self._make_client() as client:
                table = client.get_table("instance-id", "table-id")
                expected_predicate = expected_retryables.__contains__
                retry_fn_mock.side_effect = RuntimeError("stop early")
                with mock.patch(
                    "google.api_core.retry.if_exception_type"
                ) as predicate_builder_mock:
                    predicate_builder_mock.return_value = expected_predicate
                    with pytest.raises(Exception):
                        test_fn = table.__getattribute__(fn_name)
                        test_fn(*fn_args, retryable_errors=input_retryables)
                    predicate_builder_mock.assert_called_once_with(
                        *expected_retryables, *extra_retryables
                    )
                    retry_call_args = retry_fn_mock.call_args_list[0].args
                    assert retry_call_args[1] is expected_predicate

    @pytest.mark.parametrize(
        "fn_name,fn_args,gapic_fn",
        [
            ("read_rows_stream", (ReadRowsQuery(),), "read_rows"),
            ("read_rows", (ReadRowsQuery(),), "read_rows"),
            ("read_row", (b"row_key",), "read_rows"),
            ("read_rows_sharded", ([ReadRowsQuery()],), "read_rows"),
            ("row_exists", (b"row_key",), "read_rows"),
            ("sample_row_keys", (), "sample_row_keys"),
            ("mutate_row", (b"row_key", [mutations.DeleteAllFromRow()]), "mutate_row"),
            (
                "bulk_mutate_rows",
                ([mutations.RowMutationEntry(b"key", [mutations.DeleteAllFromRow()])],),
                "mutate_rows",
            ),
            ("check_and_mutate_row", (b"row_key", None), "check_and_mutate_row"),
            (
                "read_modify_write_row",
                (b"row_key", IncrementRule("f", "q")),
                "read_modify_write_row",
            ),
        ],
    )
    @pytest.mark.parametrize("include_app_profile", [True, False])
    def test_call_metadata(self, include_app_profile, fn_name, fn_args, gapic_fn):
        profile = "profile" if include_app_profile else None
        client = self._make_client()
        transport_mock = mock.MagicMock()
        rpc_mock = CrossSync._Sync_Impl.Mock()
        transport_mock._wrapped_methods.__getitem__.return_value = rpc_mock
        gapic_client = client._gapic_client
        gapic_client._transport = transport_mock
        gapic_client._is_universe_domain_valid = True
        table = self._make_one(client, app_profile_id=profile)
        try:
            test_fn = table.__getattribute__(fn_name)
            maybe_stream = test_fn(*fn_args)
            [i for i in maybe_stream]
        except Exception:
            pass
        assert rpc_mock.call_count == 1
        kwargs = rpc_mock.call_args_list[0][1]
        metadata = kwargs["metadata"]
        assert len(metadata) == 1
        assert metadata[0][0] == "x-goog-request-params"
        routing_str = metadata[0][1]
        assert self._expected_routing_header(table) in routing_str
        if include_app_profile:
            assert "app_profile_id=profile" in routing_str
        else:
            assert "app_profile_id=" in routing_str

    @staticmethod
    def _expected_routing_header(table):
        """the expected routing header for this _ApiSurface type"""
        return f"table_name={table.table_name}"


@CrossSync._Sync_Impl.add_mapping_decorator("TestAuthorizedView")
class TestAuthorizedView(CrossSync._Sync_Impl.TestTable):
    """
    Inherit tests from TestTableAsync, with some modifications
    """

    @staticmethod
    def _get_target_class():
        return CrossSync._Sync_Impl.AuthorizedView

    def _make_one(
        self,
        client,
        instance_id="instance",
        table_id="table",
        view_id="view",
        app_profile_id=None,
        **kwargs,
    ):
        return self._get_target_class()(
            client, instance_id, table_id, view_id, app_profile_id, **kwargs
        )

    @staticmethod
    def _expected_routing_header(view):
        """the expected routing header for this _ApiSurface type"""
        return f"authorized_view_name={view.authorized_view_name}"

    def test_ctor(self):
        from google.cloud.bigtable.data._helpers import _WarmedInstanceKey

        expected_table_id = "table-id"
        expected_instance_id = "instance-id"
        expected_view_id = "view_id"
        expected_app_profile_id = "app-profile-id"
        expected_operation_timeout = 123
        expected_attempt_timeout = 12
        expected_read_rows_operation_timeout = 1.5
        expected_read_rows_attempt_timeout = 0.5
        expected_mutate_rows_operation_timeout = 2.5
        expected_mutate_rows_attempt_timeout = 0.75
        client = self._make_client()
        assert not client._active_instances
        view = self._get_target_class()(
            client,
            expected_instance_id,
            expected_table_id,
            expected_view_id,
            expected_app_profile_id,
            default_operation_timeout=expected_operation_timeout,
            default_attempt_timeout=expected_attempt_timeout,
            default_read_rows_operation_timeout=expected_read_rows_operation_timeout,
            default_read_rows_attempt_timeout=expected_read_rows_attempt_timeout,
            default_mutate_rows_operation_timeout=expected_mutate_rows_operation_timeout,
            default_mutate_rows_attempt_timeout=expected_mutate_rows_attempt_timeout,
        )
        CrossSync._Sync_Impl.yield_to_event_loop()
        assert view.table_id == expected_table_id
        assert (
            view.table_name
            == f"projects/{client.project}/instances/{expected_instance_id}/tables/{expected_table_id}"
        )
        assert view.instance_id == expected_instance_id
        assert (
            view.instance_name
            == f"projects/{client.project}/instances/{expected_instance_id}"
        )
        assert view.authorized_view_id == expected_view_id
        assert (
            view.authorized_view_name
            == f"projects/{client.project}/instances/{expected_instance_id}/tables/{expected_table_id}/authorizedViews/{expected_view_id}"
        )
        assert view.app_profile_id == expected_app_profile_id
        assert view.client is client
        instance_key = _WarmedInstanceKey(view.instance_name, view.app_profile_id)
        assert instance_key in client._active_instances
        assert client._instance_owners[instance_key] == {id(view)}
        assert view.default_operation_timeout == expected_operation_timeout
        assert view.default_attempt_timeout == expected_attempt_timeout
        assert (
            view.default_read_rows_operation_timeout
            == expected_read_rows_operation_timeout
        )
        assert (
            view.default_read_rows_attempt_timeout == expected_read_rows_attempt_timeout
        )
        assert (
            view.default_mutate_rows_operation_timeout
            == expected_mutate_rows_operation_timeout
        )
        assert (
            view.default_mutate_rows_attempt_timeout
            == expected_mutate_rows_attempt_timeout
        )
        view._register_instance_future
        assert view._register_instance_future.done()
        assert not view._register_instance_future.cancelled()
        assert view._register_instance_future.exception() is None
        client.close()


@CrossSync._Sync_Impl.add_mapping_decorator("TestReadRows")
class TestReadRows:
    """
    Tests for table.read_rows and related methods.
    """

    @staticmethod
    def _get_operation_class():
        return CrossSync._Sync_Impl._ReadRowsOperation

    def _make_client(self, *args, **kwargs):
        return CrossSync._Sync_Impl.TestBigtableDataClient._make_client(*args, **kwargs)

    def _make_table(self, *args, **kwargs):
        client_mock = mock.Mock()
        client_mock._register_instance.side_effect = (
            lambda *args, **kwargs: CrossSync._Sync_Impl.yield_to_event_loop()
        )
        client_mock._remove_instance_registration.side_effect = (
            lambda *args, **kwargs: CrossSync._Sync_Impl.yield_to_event_loop()
        )
        kwargs["instance_id"] = kwargs.get(
            "instance_id", args[0] if args else "instance"
        )
        kwargs["table_id"] = kwargs.get(
            "table_id", args[1] if len(args) > 1 else "table"
        )
        client_mock._gapic_client.table_path.return_value = kwargs["table_id"]
        client_mock._gapic_client.instance_path.return_value = kwargs["instance_id"]
        return CrossSync._Sync_Impl.TestTable._get_target_class()(
            client_mock, *args, **kwargs
        )

    def _make_stats(self):
        from google.cloud.bigtable_v2.types import RequestStats
        from google.cloud.bigtable_v2.types import FullReadStatsView
        from google.cloud.bigtable_v2.types import ReadIterationStats

        return RequestStats(
            full_read_stats_view=FullReadStatsView(
                read_iteration_stats=ReadIterationStats(
                    rows_seen_count=1,
                    rows_returned_count=2,
                    cells_seen_count=3,
                    cells_returned_count=4,
                )
            )
        )

    @staticmethod
    def _make_chunk(*args, **kwargs):
        from google.cloud.bigtable_v2 import ReadRowsResponse

        kwargs["row_key"] = kwargs.get("row_key", b"row_key")
        kwargs["family_name"] = kwargs.get("family_name", "family_name")
        kwargs["qualifier"] = kwargs.get("qualifier", b"qualifier")
        kwargs["value"] = kwargs.get("value", b"value")
        kwargs["commit_row"] = kwargs.get("commit_row", True)
        return ReadRowsResponse.CellChunk(*args, **kwargs)

    @staticmethod
    def _make_gapic_stream(
        chunk_list: list[ReadRowsResponse.CellChunk | Exception], sleep_time=0
    ):
        from google.cloud.bigtable_v2 import ReadRowsResponse

        class mock_stream:
            def __init__(self, chunk_list, sleep_time):
                self.chunk_list = chunk_list
                self.idx = -1
                self.sleep_time = sleep_time

            def __iter__(self):
                return self

            def __next__(self):
                self.idx += 1
                if len(self.chunk_list) > self.idx:
                    if sleep_time:
                        CrossSync._Sync_Impl.sleep(self.sleep_time)
                    chunk = self.chunk_list[self.idx]
                    if isinstance(chunk, Exception):
                        raise chunk
                    else:
                        return ReadRowsResponse(chunks=[chunk])
                raise CrossSync._Sync_Impl.StopIteration

            def cancel(self):
                pass

        return mock_stream(chunk_list, sleep_time)

    def execute_fn(self, table, *args, **kwargs):
        return table.read_rows(*args, **kwargs)

    def test_read_rows(self):
        query = ReadRowsQuery()
        chunks = [
            self._make_chunk(row_key=b"test_1"),
            self._make_chunk(row_key=b"test_2"),
        ]
        with self._make_table() as table:
            read_rows = table.client._gapic_client.read_rows
            read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream(
                chunks
            )
            results = self.execute_fn(table, query, operation_timeout=3)
            assert len(results) == 2
            assert results[0].row_key == b"test_1"
            assert results[1].row_key == b"test_2"

    def test_read_rows_stream(self):
        query = ReadRowsQuery()
        chunks = [
            self._make_chunk(row_key=b"test_1"),
            self._make_chunk(row_key=b"test_2"),
        ]
        with self._make_table() as table:
            read_rows = table.client._gapic_client.read_rows
            read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream(
                chunks
            )
            gen = table.read_rows_stream(query, operation_timeout=3)
            results = [row for row in gen]
            assert len(results) == 2
            assert results[0].row_key == b"test_1"
            assert results[1].row_key == b"test_2"

    @pytest.mark.parametrize("include_app_profile", [True, False])
    def test_read_rows_query_matches_request(self, include_app_profile):
        from google.cloud.bigtable.data import RowRange
        from google.cloud.bigtable.data.row_filters import PassAllFilter

        app_profile_id = "app_profile_id" if include_app_profile else None
        with self._make_table(app_profile_id=app_profile_id) as table:
            read_rows = table.client._gapic_client.read_rows
            read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream([])
            row_keys = [b"test_1", "test_2"]
            row_ranges = RowRange("1start", "2end")
            filter_ = PassAllFilter(True)
            limit = 99
            query = ReadRowsQuery(
                row_keys=row_keys,
                row_ranges=row_ranges,
                row_filter=filter_,
                limit=limit,
            )
            results = table.read_rows(query, operation_timeout=3)
            assert len(results) == 0
            call_request = read_rows.call_args_list[0][0][0]
            query_pb = query._to_pb(table)
            assert call_request == query_pb

    @pytest.mark.parametrize("operation_timeout", [0.001, 0.023, 0.1])
    def test_read_rows_timeout(self, operation_timeout):
        with self._make_table() as table:
            read_rows = table.client._gapic_client.read_rows
            query = ReadRowsQuery()
            chunks = [self._make_chunk(row_key=b"test_1")]
            read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream(
                chunks, sleep_time=0.15
            )
            try:
                table.read_rows(query, operation_timeout=operation_timeout)
            except core_exceptions.DeadlineExceeded as e:
                assert (
                    e.message
                    == f"operation_timeout of {operation_timeout:0.1f}s exceeded"
                )

    @pytest.mark.parametrize(
        "per_request_t, operation_t, expected_num",
        [(0.05, 0.08, 2), (0.05, 0.14, 3), (0.05, 0.24, 5)],
    )
    def test_read_rows_attempt_timeout(self, per_request_t, operation_t, expected_num):
        """Ensures that the attempt_timeout is respected and that the number of
        requests is as expected.

        operation_timeout does not cancel the request, so we expect the number of
        requests to be the ceiling of operation_timeout / attempt_timeout."""
        from google.cloud.bigtable.data.exceptions import RetryExceptionGroup

        expected_last_timeout = operation_t - (expected_num - 1) * per_request_t
        with mock.patch("random.uniform", side_effect=lambda a, b: 0):
            with self._make_table() as table:
                read_rows = table.client._gapic_client.read_rows
                read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream(
                    chunks, sleep_time=per_request_t
                )
                query = ReadRowsQuery()
                chunks = [core_exceptions.DeadlineExceeded("mock deadline")]
                try:
                    table.read_rows(
                        query,
                        operation_timeout=operation_t,
                        attempt_timeout=per_request_t,
                    )
                except core_exceptions.DeadlineExceeded as e:
                    retry_exc = e.__cause__
                    if expected_num == 0:
                        assert retry_exc is None
                    else:
                        assert type(retry_exc) is RetryExceptionGroup
                        assert f"{expected_num} failed attempts" in str(retry_exc)
                        assert len(retry_exc.exceptions) == expected_num
                        for sub_exc in retry_exc.exceptions:
                            assert sub_exc.message == "mock deadline"
                assert read_rows.call_count == expected_num
                for _, call_kwargs in read_rows.call_args_list[:-1]:
                    assert call_kwargs["timeout"] == per_request_t
                    assert call_kwargs["retry"] is None
                assert (
                    abs(
                        read_rows.call_args_list[-1][1]["timeout"]
                        - expected_last_timeout
                    )
                    < 0.05
                )

    @pytest.mark.parametrize(
        "exc_type",
        [
            core_exceptions.Aborted,
            core_exceptions.DeadlineExceeded,
            core_exceptions.ServiceUnavailable,
        ],
    )
    def test_read_rows_retryable_error(self, exc_type):
        with self._make_table() as table:
            read_rows = table.client._gapic_client.read_rows
            read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream(
                [expected_error]
            )
            query = ReadRowsQuery()
            expected_error = exc_type("mock error")
            try:
                table.read_rows(query, operation_timeout=0.1)
            except core_exceptions.DeadlineExceeded as e:
                retry_exc = e.__cause__
                root_cause = retry_exc.exceptions[0]
                assert type(root_cause) is exc_type
                assert root_cause == expected_error

    @pytest.mark.parametrize(
        "exc_type",
        [
            core_exceptions.Cancelled,
            core_exceptions.PreconditionFailed,
            core_exceptions.NotFound,
            core_exceptions.PermissionDenied,
            core_exceptions.Conflict,
            core_exceptions.InternalServerError,
            core_exceptions.TooManyRequests,
            core_exceptions.ResourceExhausted,
            InvalidChunk,
        ],
    )
    def test_read_rows_non_retryable_error(self, exc_type):
        with self._make_table() as table:
            read_rows = table.client._gapic_client.read_rows
            read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream(
                [expected_error]
            )
            query = ReadRowsQuery()
            expected_error = exc_type("mock error")
            try:
                table.read_rows(query, operation_timeout=0.1)
            except exc_type as e:
                assert e == expected_error

    def test_read_rows_revise_request(self):
        """Ensure that _revise_request is called between retries"""
        from google.cloud.bigtable.data.exceptions import InvalidChunk
        from google.cloud.bigtable_v2.types import RowSet

        return_val = RowSet()
        with mock.patch.object(
            self._get_operation_class(), "_revise_request_rowset"
        ) as revise_rowset:
            revise_rowset.return_value = return_val
            with self._make_table() as table:
                read_rows = table.client._gapic_client.read_rows
                read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream(
                    chunks
                )
                row_keys = [b"test_1", b"test_2", b"test_3"]
                query = ReadRowsQuery(row_keys=row_keys)
                chunks = [
                    self._make_chunk(row_key=b"test_1"),
                    core_exceptions.Aborted("mock retryable error"),
                ]
                try:
                    table.read_rows(query)
                except InvalidChunk:
                    revise_rowset.assert_called()
                    first_call_kwargs = revise_rowset.call_args_list[0].kwargs
                    assert first_call_kwargs["row_set"] == query._to_pb(table).rows
                    assert first_call_kwargs["last_seen_row_key"] == b"test_1"
                    revised_call = read_rows.call_args_list[1].args[0]
                    assert revised_call.rows == return_val

    def test_read_rows_default_timeouts(self):
        """Ensure that the default timeouts are set on the read rows operation when not overridden"""
        operation_timeout = 8
        attempt_timeout = 4
        with mock.patch.object(self._get_operation_class(), "__init__") as mock_op:
            mock_op.side_effect = RuntimeError("mock error")
            with self._make_table(
                default_read_rows_operation_timeout=operation_timeout,
                default_read_rows_attempt_timeout=attempt_timeout,
            ) as table:
                try:
                    table.read_rows(ReadRowsQuery())
                except RuntimeError:
                    pass
                kwargs = mock_op.call_args_list[0].kwargs
                assert kwargs["operation_timeout"] == operation_timeout
                assert kwargs["attempt_timeout"] == attempt_timeout

    def test_read_rows_default_timeout_override(self):
        """When timeouts are passed, they overwrite default values"""
        operation_timeout = 8
        attempt_timeout = 4
        with mock.patch.object(self._get_operation_class(), "__init__") as mock_op:
            mock_op.side_effect = RuntimeError("mock error")
            with self._make_table(
                default_operation_timeout=99, default_attempt_timeout=97
            ) as table:
                try:
                    table.read_rows(
                        ReadRowsQuery(),
                        operation_timeout=operation_timeout,
                        attempt_timeout=attempt_timeout,
                    )
                except RuntimeError:
                    pass
                kwargs = mock_op.call_args_list[0].kwargs
                assert kwargs["operation_timeout"] == operation_timeout
                assert kwargs["attempt_timeout"] == attempt_timeout

    def test_read_row(self):
        """Test reading a single row"""
        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:
                expected_result = object()
                read_rows.side_effect = lambda *args, **kwargs: [expected_result]
                expected_op_timeout = 8
                expected_req_timeout = 4
                row = table.read_row(
                    row_key,
                    operation_timeout=expected_op_timeout,
                    attempt_timeout=expected_req_timeout,
                )
                assert row == expected_result
                assert read_rows.call_count == 1
                (args, kwargs) = read_rows.call_args_list[0]
                assert kwargs["operation_timeout"] == expected_op_timeout
                assert kwargs["attempt_timeout"] == expected_req_timeout
                assert len(args) == 1
                assert isinstance(args[0], ReadRowsQuery)
                query = args[0]
                assert query.row_keys == [row_key]
                assert query.row_ranges == []
                assert query.limit == 1

    def test_read_row_w_filter(self):
        """Test reading a single row with an added filter"""
        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:
                expected_result = object()
                read_rows.side_effect = lambda *args, **kwargs: [expected_result]
                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
                row = 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 kwargs["operation_timeout"] == expected_op_timeout
                assert kwargs["attempt_timeout"] == expected_req_timeout
                assert len(args) == 1
                assert isinstance(args[0], ReadRowsQuery)
                query = args[0]
                assert query.row_keys == [row_key]
                assert query.row_ranges == []
                assert query.limit == 1
                assert query.filter == expected_filter

    def test_read_row_no_response(self):
        """should return None if row does not exist"""
        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:
                read_rows.side_effect = lambda *args, **kwargs: []
                expected_op_timeout = 8
                expected_req_timeout = 4
                result = table.read_row(
                    row_key,
                    operation_timeout=expected_op_timeout,
                    attempt_timeout=expected_req_timeout,
                )
                assert result is None
                assert read_rows.call_count == 1
                (args, kwargs) = read_rows.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 query.row_keys == [row_key]
                assert query.row_ranges == []
                assert query.limit == 1

    @pytest.mark.parametrize(
        "return_value,expected_result",
        [([], False), ([object()], True), ([object(), object()], True)],
    )
    def test_row_exists(self, return_value, expected_result):
        """Test checking for row existence"""
        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:
                read_rows.side_effect = lambda *args, **kwargs: return_value
                expected_op_timeout = 1
                expected_req_timeout = 2
                result = 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 kwargs["operation_timeout"] == expected_op_timeout
                assert kwargs["attempt_timeout"] == expected_req_timeout
                assert isinstance(args[0], ReadRowsQuery)
                expected_filter = {
                    "chain": {
                        "filters": [
                            {"cells_per_row_limit_filter": 1},
                            {"strip_value_transformer": True},
                        ]
                    }
                }
                query = args[0]
                assert query.row_keys == [row_key]
                assert query.row_ranges == []
                assert query.limit == 1
                assert query.filter._to_dict() == expected_filter


class TestReadRowsSharded:
    def _make_client(self, *args, **kwargs):
        return CrossSync._Sync_Impl.TestBigtableDataClient._make_client(*args, **kwargs)

    def test_read_rows_sharded_empty_query(self):
        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with pytest.raises(ValueError) as exc:
                    table.read_rows_sharded([])
                assert "empty sharded_query" in str(exc.value)

    def test_read_rows_sharded_multiple_queries(self):
        """Test with multiple queries. Should return results from both"""
        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    table.client._gapic_client, "read_rows"
                ) as read_rows:
                    read_rows.side_effect = lambda *args, **kwargs: CrossSync._Sync_Impl.TestReadRows._make_gapic_stream(
                        [
                            CrossSync._Sync_Impl.TestReadRows._make_chunk(row_key=k)
                            for k in args[0].rows.row_keys
                        ]
                    )
                    query_1 = ReadRowsQuery(b"test_1")
                    query_2 = ReadRowsQuery(b"test_2")
                    result = table.read_rows_sharded([query_1, query_2])
                    assert len(result) == 2
                    assert result[0].row_key == b"test_1"
                    assert result[1].row_key == b"test_2"

    @pytest.mark.parametrize("n_queries", [1, 2, 5, 11, 24])
    def test_read_rows_sharded_multiple_queries_calls(self, n_queries):
        """Each query should trigger a separate read_rows call"""
        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(table, "read_rows") as read_rows:
                    query_list = [ReadRowsQuery() for _ in range(n_queries)]
                    table.read_rows_sharded(query_list)
                    assert read_rows.call_count == n_queries

    def test_read_rows_sharded_errors(self):
        """Errors should be exposed as ShardedReadRowsExceptionGroups"""
        from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup
        from google.cloud.bigtable.data.exceptions import FailedQueryShardError

        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(table, "read_rows") as read_rows:
                    read_rows.side_effect = RuntimeError("mock error")
                    query_1 = ReadRowsQuery(b"test_1")
                    query_2 = ReadRowsQuery(b"test_2")
                    with pytest.raises(ShardedReadRowsExceptionGroup) as exc:
                        table.read_rows_sharded([query_1, query_2])
                    exc_group = exc.value
                    assert isinstance(exc_group, ShardedReadRowsExceptionGroup)
                    assert len(exc.value.exceptions) == 2
                    assert isinstance(exc.value.exceptions[0], FailedQueryShardError)
                    assert isinstance(exc.value.exceptions[0].__cause__, RuntimeError)
                    assert exc.value.exceptions[0].index == 0
                    assert exc.value.exceptions[0].query == query_1
                    assert isinstance(exc.value.exceptions[1], FailedQueryShardError)
                    assert isinstance(exc.value.exceptions[1].__cause__, RuntimeError)
                    assert exc.value.exceptions[1].index == 1
                    assert exc.value.exceptions[1].query == query_2

    def test_read_rows_sharded_concurrent(self):
        """Ensure sharded requests are concurrent"""
        import time

        def mock_call(*args, **kwargs):
            CrossSync._Sync_Impl.sleep(0.1)
            return [mock.Mock()]

        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(table, "read_rows") as read_rows:
                    read_rows.side_effect = mock_call
                    queries = [ReadRowsQuery() for _ in range(10)]
                    start_time = time.monotonic()
                    result = table.read_rows_sharded(queries)
                    call_time = time.monotonic() - start_time
                    assert read_rows.call_count == 10
                    assert len(result) == 10
                    assert call_time < 0.5

    def test_read_rows_sharded_concurrency_limit(self):
        """Only 10 queries should be processed concurrently. Others should be queued

        Should start a new query as soon as previous finishes"""
        from google.cloud.bigtable.data._helpers import _CONCURRENCY_LIMIT

        assert _CONCURRENCY_LIMIT == 10
        num_queries = 15
        increment_time = 0.05
        max_time = increment_time * (_CONCURRENCY_LIMIT - 1)
        rpc_times = [min(i * increment_time, max_time) for i in range(num_queries)]

        def mock_call(*args, **kwargs):
            next_sleep = rpc_times.pop(0)
            asyncio.sleep(next_sleep)
            return [mock.Mock()]

        starting_timeout = 10
        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(table, "read_rows") as read_rows:
                    read_rows.side_effect = mock_call
                    queries = [ReadRowsQuery() for _ in range(num_queries)]
                    table.read_rows_sharded(queries, operation_timeout=starting_timeout)
                    assert read_rows.call_count == num_queries
                    rpc_start_list = [
                        starting_timeout - kwargs["operation_timeout"]
                        for (_, kwargs) in read_rows.call_args_list
                    ]
                    eps = 0.01
                    assert all(
                        (rpc_start_list[i] < eps for i in range(_CONCURRENCY_LIMIT))
                    )
                    for i in range(num_queries - _CONCURRENCY_LIMIT):
                        idx = i + _CONCURRENCY_LIMIT
                        assert rpc_start_list[idx] - i * increment_time < eps

    def test_read_rows_sharded_expirary(self):
        """If the operation times out before all shards complete, should raise
        a ShardedReadRowsExceptionGroup"""
        from google.cloud.bigtable.data._helpers import _CONCURRENCY_LIMIT
        from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup
        from google.api_core.exceptions import DeadlineExceeded

        operation_timeout = 0.1
        num_queries = 15
        sleeps = [0] * _CONCURRENCY_LIMIT + [DeadlineExceeded("times up")] * (
            num_queries - _CONCURRENCY_LIMIT
        )

        def mock_call(*args, **kwargs):
            next_item = sleeps.pop(0)
            if isinstance(next_item, Exception):
                raise next_item
            else:
                asyncio.sleep(next_item)
            return [mock.Mock()]

        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(table, "read_rows") as read_rows:
                    read_rows.side_effect = mock_call
                    queries = [ReadRowsQuery() for _ in range(num_queries)]
                    with pytest.raises(ShardedReadRowsExceptionGroup) as exc:
                        table.read_rows_sharded(
                            queries, operation_timeout=operation_timeout
                        )
                    assert isinstance(exc.value, ShardedReadRowsExceptionGroup)
                    assert len(exc.value.exceptions) == num_queries - _CONCURRENCY_LIMIT
                    assert len(exc.value.successful_rows) == _CONCURRENCY_LIMIT

    def test_read_rows_sharded_negative_batch_timeout(self):
        """try to run with batch that starts after operation timeout

        They should raise DeadlineExceeded errors"""
        from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup
        from google.cloud.bigtable.data._helpers import _CONCURRENCY_LIMIT
        from google.api_core.exceptions import DeadlineExceeded

        def mock_call(*args, **kwargs):
            CrossSync._Sync_Impl.sleep(0.06)
            return [mock.Mock()]

        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(table, "read_rows") as read_rows:
                    read_rows.side_effect = mock_call
                    num_calls = 15
                    queries = [ReadRowsQuery() for _ in range(num_calls)]
                    with pytest.raises(ShardedReadRowsExceptionGroup) as exc:
                        table.read_rows_sharded(queries, operation_timeout=0.05)
                    assert isinstance(exc.value, ShardedReadRowsExceptionGroup)
                    assert len(exc.value.exceptions) >= num_calls - _CONCURRENCY_LIMIT
                    assert all(
                        (
                            isinstance(e.__cause__, DeadlineExceeded)
                            for e in exc.value.exceptions
                        )
                    )


class TestSampleRowKeys:
    def _make_client(self, *args, **kwargs):
        return CrossSync._Sync_Impl.TestBigtableDataClient._make_client(*args, **kwargs)

    def _make_gapic_stream(self, sample_list: list[tuple[bytes, int]]):
        from google.cloud.bigtable_v2.types import SampleRowKeysResponse

        for value in sample_list:
            yield SampleRowKeysResponse(row_key=value[0], offset_bytes=value[1])

    def test_sample_row_keys(self):
        """Test that method returns the expected key samples"""
        samples = [(b"test_1", 0), (b"test_2", 100), (b"test_3", 200)]
        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    table.client._gapic_client,
                    "sample_row_keys",
                    CrossSync._Sync_Impl.Mock(),
                ) as sample_row_keys:
                    sample_row_keys.return_value = self._make_gapic_stream(samples)
                    result = table.sample_row_keys()
                    assert len(result) == 3
                    assert all((isinstance(r, tuple) for r in result))
                    assert all((isinstance(r[0], bytes) for r in result))
                    assert all((isinstance(r[1], int) for r in result))
                    assert result[0] == samples[0]
                    assert result[1] == samples[1]
                    assert result[2] == samples[2]

    def test_sample_row_keys_bad_timeout(self):
        """should raise error if timeout is negative"""
        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with pytest.raises(ValueError) as e:
                    table.sample_row_keys(operation_timeout=-1)
                    assert "operation_timeout must be greater than 0" in str(e.value)
                with pytest.raises(ValueError) as e:
                    table.sample_row_keys(attempt_timeout=-1)
                    assert "attempt_timeout must be greater than 0" in str(e.value)

    def test_sample_row_keys_default_timeout(self):
        """Should fallback to using table default operation_timeout"""
        expected_timeout = 99
        with self._make_client() as client:
            with client.get_table(
                "i",
                "t",
                default_operation_timeout=expected_timeout,
                default_attempt_timeout=expected_timeout,
            ) as table:
                with mock.patch.object(
                    table.client._gapic_client,
                    "sample_row_keys",
                    CrossSync._Sync_Impl.Mock(),
                ) as sample_row_keys:
                    sample_row_keys.return_value = self._make_gapic_stream([])
                    result = table.sample_row_keys()
                    (_, kwargs) = sample_row_keys.call_args
                    assert abs(kwargs["timeout"] - expected_timeout) < 0.1
                    assert result == []
                    assert kwargs["retry"] is None

    def test_sample_row_keys_gapic_params(self):
        """make sure arguments are propagated to gapic call as expected"""
        expected_timeout = 10
        expected_profile = "test1"
        instance = "instance_name"
        table_id = "my_table"
        with self._make_client() as client:
            with client.get_table(
                instance, table_id, app_profile_id=expected_profile
            ) as table:
                with mock.patch.object(
                    table.client._gapic_client,
                    "sample_row_keys",
                    CrossSync._Sync_Impl.Mock(),
                ) as sample_row_keys:
                    sample_row_keys.return_value = self._make_gapic_stream([])
                    table.sample_row_keys(attempt_timeout=expected_timeout)
                    (args, kwargs) = sample_row_keys.call_args
                    assert len(args) == 0
                    assert len(kwargs) == 3
                    assert kwargs["timeout"] == expected_timeout
                    assert kwargs["retry"] is None
                    request = kwargs["request"]
                    assert request.app_profile_id == expected_profile
                    assert request.table_name == table.table_name

    @pytest.mark.parametrize(
        "retryable_exception",
        [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable],
    )
    def test_sample_row_keys_retryable_errors(self, retryable_exception):
        """retryable errors should be retried until timeout"""
        from google.api_core.exceptions import DeadlineExceeded
        from google.cloud.bigtable.data.exceptions import RetryExceptionGroup

        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    table.client._gapic_client,
                    "sample_row_keys",
                    CrossSync._Sync_Impl.Mock(),
                ) as sample_row_keys:
                    sample_row_keys.side_effect = retryable_exception("mock")
                    with pytest.raises(DeadlineExceeded) as e:
                        table.sample_row_keys(operation_timeout=0.05)
                    cause = e.value.__cause__
                    assert isinstance(cause, RetryExceptionGroup)
                    assert len(cause.exceptions) > 0
                    assert isinstance(cause.exceptions[0], retryable_exception)

    @pytest.mark.parametrize(
        "non_retryable_exception",
        [
            core_exceptions.OutOfRange,
            core_exceptions.NotFound,
            core_exceptions.FailedPrecondition,
            RuntimeError,
            ValueError,
            core_exceptions.Aborted,
        ],
    )
    def test_sample_row_keys_non_retryable_errors(self, non_retryable_exception):
        """non-retryable errors should cause a raise"""
        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    table.client._gapic_client,
                    "sample_row_keys",
                    CrossSync._Sync_Impl.Mock(),
                ) as sample_row_keys:
                    sample_row_keys.side_effect = non_retryable_exception("mock")
                    with pytest.raises(non_retryable_exception):
                        table.sample_row_keys()


class TestMutateRow:
    def _make_client(self, *args, **kwargs):
        return CrossSync._Sync_Impl.TestBigtableDataClient._make_client(*args, **kwargs)

    @pytest.mark.parametrize(
        "mutation_arg",
        [
            mutations.SetCell("family", b"qualifier", b"value"),
            mutations.SetCell(
                "family", b"qualifier", b"value", timestamp_micros=1234567890
            ),
            mutations.DeleteRangeFromColumn("family", b"qualifier"),
            mutations.DeleteAllFromFamily("family"),
            mutations.DeleteAllFromRow(),
            [mutations.SetCell("family", b"qualifier", b"value")],
            [
                mutations.DeleteRangeFromColumn("family", b"qualifier"),
                mutations.DeleteAllFromRow(),
            ],
        ],
    )
    def test_mutate_row(self, mutation_arg):
        """Test mutations with no errors"""
        expected_attempt_timeout = 19
        with self._make_client(project="project") as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "mutate_row"
                ) as mock_gapic:
                    mock_gapic.return_value = None
                    table.mutate_row(
                        "row_key",
                        mutation_arg,
                        attempt_timeout=expected_attempt_timeout,
                    )
                    assert mock_gapic.call_count == 1
                    kwargs = mock_gapic.call_args_list[0].kwargs
                    request = kwargs["request"]
                    assert (
                        request.table_name
                        == "projects/project/instances/instance/tables/table"
                    )
                    assert request.row_key == b"row_key"
                    formatted_mutations = (
                        [mutation._to_pb() for mutation in mutation_arg]
                        if isinstance(mutation_arg, list)
                        else [mutation_arg._to_pb()]
                    )
                    assert request.mutations == formatted_mutations
                    assert kwargs["timeout"] == expected_attempt_timeout
                    assert kwargs["retry"] is None

    @pytest.mark.parametrize(
        "retryable_exception",
        [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable],
    )
    def test_mutate_row_retryable_errors(self, retryable_exception):
        from google.api_core.exceptions import DeadlineExceeded
        from google.cloud.bigtable.data.exceptions import RetryExceptionGroup

        with self._make_client(project="project") as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "mutate_row"
                ) as mock_gapic:
                    mock_gapic.side_effect = retryable_exception("mock")
                    with pytest.raises(DeadlineExceeded) as e:
                        mutation = mutations.DeleteAllFromRow()
                        assert mutation.is_idempotent() is True
                        table.mutate_row("row_key", mutation, operation_timeout=0.01)
                    cause = e.value.__cause__
                    assert isinstance(cause, RetryExceptionGroup)
                    assert isinstance(cause.exceptions[0], retryable_exception)

    @pytest.mark.parametrize(
        "retryable_exception",
        [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable],
    )
    def test_mutate_row_non_idempotent_retryable_errors(self, retryable_exception):
        """Non-idempotent mutations should not be retried"""
        with self._make_client(project="project") as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "mutate_row"
                ) as mock_gapic:
                    mock_gapic.side_effect = retryable_exception("mock")
                    with pytest.raises(retryable_exception):
                        mutation = mutations.SetCell(
                            "family", b"qualifier", b"value", -1
                        )
                        assert mutation.is_idempotent() is False
                        table.mutate_row("row_key", mutation, operation_timeout=0.2)

    @pytest.mark.parametrize(
        "non_retryable_exception",
        [
            core_exceptions.OutOfRange,
            core_exceptions.NotFound,
            core_exceptions.FailedPrecondition,
            RuntimeError,
            ValueError,
            core_exceptions.Aborted,
        ],
    )
    def test_mutate_row_non_retryable_errors(self, non_retryable_exception):
        with self._make_client(project="project") as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "mutate_row"
                ) as mock_gapic:
                    mock_gapic.side_effect = non_retryable_exception("mock")
                    with pytest.raises(non_retryable_exception):
                        mutation = mutations.SetCell(
                            "family",
                            b"qualifier",
                            b"value",
                            timestamp_micros=1234567890,
                        )
                        assert mutation.is_idempotent() is True
                        table.mutate_row("row_key", mutation, operation_timeout=0.2)

    @pytest.mark.parametrize("mutations", [[], None])
    def test_mutate_row_no_mutations(self, mutations):
        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with pytest.raises(ValueError) as e:
                    table.mutate_row("key", mutations=mutations)
                    assert e.value.args[0] == "No mutations provided"


class TestBulkMutateRows:
    def _make_client(self, *args, **kwargs):
        return CrossSync._Sync_Impl.TestBigtableDataClient._make_client(*args, **kwargs)

    def _mock_response(self, response_list):
        from google.cloud.bigtable_v2.types import MutateRowsResponse
        from google.rpc import status_pb2

        statuses = []
        for response in response_list:
            if isinstance(response, core_exceptions.GoogleAPICallError):
                statuses.append(
                    status_pb2.Status(
                        message=str(response), code=response.grpc_status_code.value[0]
                    )
                )
            else:
                statuses.append(status_pb2.Status(code=0))
        entries = [
            MutateRowsResponse.Entry(index=i, status=statuses[i])
            for i in range(len(response_list))
        ]

        def generator():
            yield MutateRowsResponse(entries=entries)

        return generator()

    @pytest.mark.parametrize(
        "mutation_arg",
        [
            [mutations.SetCell("family", b"qualifier", b"value")],
            [
                mutations.SetCell(
                    "family", b"qualifier", b"value", timestamp_micros=1234567890
                )
            ],
            [mutations.DeleteRangeFromColumn("family", b"qualifier")],
            [mutations.DeleteAllFromFamily("family")],
            [mutations.DeleteAllFromRow()],
            [mutations.SetCell("family", b"qualifier", b"value")],
            [
                mutations.DeleteRangeFromColumn("family", b"qualifier"),
                mutations.DeleteAllFromRow(),
            ],
        ],
    )
    def test_bulk_mutate_rows(self, mutation_arg):
        """Test mutations with no errors"""
        expected_attempt_timeout = 19
        with self._make_client(project="project") as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "mutate_rows"
                ) as mock_gapic:
                    mock_gapic.return_value = self._mock_response([None])
                    bulk_mutation = mutations.RowMutationEntry(b"row_key", mutation_arg)
                    table.bulk_mutate_rows(
                        [bulk_mutation], attempt_timeout=expected_attempt_timeout
                    )
                    assert mock_gapic.call_count == 1
                    kwargs = mock_gapic.call_args[1]
                    request = kwargs["request"]
                    assert (
                        request.table_name
                        == "projects/project/instances/instance/tables/table"
                    )
                    assert request.entries == [bulk_mutation._to_pb()]
                    assert kwargs["timeout"] == expected_attempt_timeout
                    assert kwargs["retry"] is None

    def test_bulk_mutate_rows_multiple_entries(self):
        """Test mutations with no errors"""
        with self._make_client(project="project") as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "mutate_rows"
                ) as mock_gapic:
                    mock_gapic.return_value = self._mock_response([None, None])
                    mutation_list = [mutations.DeleteAllFromRow()]
                    entry_1 = mutations.RowMutationEntry(b"row_key_1", mutation_list)
                    entry_2 = mutations.RowMutationEntry(b"row_key_2", mutation_list)
                    table.bulk_mutate_rows([entry_1, entry_2])
                    assert mock_gapic.call_count == 1
                    kwargs = mock_gapic.call_args[1]
                    request = kwargs["request"]
                    assert (
                        request.table_name
                        == "projects/project/instances/instance/tables/table"
                    )
                    assert request.entries[0] == entry_1._to_pb()
                    assert request.entries[1] == entry_2._to_pb()

    @pytest.mark.parametrize(
        "exception",
        [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable],
    )
    def test_bulk_mutate_rows_idempotent_mutation_error_retryable(self, exception):
        """Individual idempotent mutations should be retried if they fail with a retryable error"""
        from google.cloud.bigtable.data.exceptions import (
            RetryExceptionGroup,
            FailedMutationEntryError,
            MutationsExceptionGroup,
        )

        with self._make_client(project="project") as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "mutate_rows"
                ) as mock_gapic:
                    mock_gapic.side_effect = lambda *a, **k: self._mock_response(
                        [exception("mock")]
                    )
                    with pytest.raises(MutationsExceptionGroup) as e:
                        mutation = mutations.DeleteAllFromRow()
                        entry = mutations.RowMutationEntry(b"row_key", [mutation])
                        assert mutation.is_idempotent() is True
                        table.bulk_mutate_rows([entry], operation_timeout=0.05)
                    assert len(e.value.exceptions) == 1
                    failed_exception = e.value.exceptions[0]
                    assert "non-idempotent" not in str(failed_exception)
                    assert isinstance(failed_exception, FailedMutationEntryError)
                    cause = failed_exception.__cause__
                    assert isinstance(cause, RetryExceptionGroup)
                    assert isinstance(cause.exceptions[0], exception)
                    assert isinstance(
                        cause.exceptions[-1], core_exceptions.DeadlineExceeded
                    )

    @pytest.mark.parametrize(
        "exception",
        [
            core_exceptions.OutOfRange,
            core_exceptions.NotFound,
            core_exceptions.FailedPrecondition,
            core_exceptions.Aborted,
        ],
    )
    def test_bulk_mutate_rows_idempotent_mutation_error_non_retryable(self, exception):
        """Individual idempotent mutations should not be retried if they fail with a non-retryable error"""
        from google.cloud.bigtable.data.exceptions import (
            FailedMutationEntryError,
            MutationsExceptionGroup,
        )

        with self._make_client(project="project") as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "mutate_rows"
                ) as mock_gapic:
                    mock_gapic.side_effect = lambda *a, **k: self._mock_response(
                        [exception("mock")]
                    )
                    with pytest.raises(MutationsExceptionGroup) as e:
                        mutation = mutations.DeleteAllFromRow()
                        entry = mutations.RowMutationEntry(b"row_key", [mutation])
                        assert mutation.is_idempotent() is True
                        table.bulk_mutate_rows([entry], operation_timeout=0.05)
                    assert len(e.value.exceptions) == 1
                    failed_exception = e.value.exceptions[0]
                    assert "non-idempotent" not in str(failed_exception)
                    assert isinstance(failed_exception, FailedMutationEntryError)
                    cause = failed_exception.__cause__
                    assert isinstance(cause, exception)

    @pytest.mark.parametrize(
        "retryable_exception",
        [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable],
    )
    def test_bulk_mutate_idempotent_retryable_request_errors(self, retryable_exception):
        """Individual idempotent mutations should be retried if the request fails with a retryable error"""
        from google.cloud.bigtable.data.exceptions import (
            RetryExceptionGroup,
            FailedMutationEntryError,
            MutationsExceptionGroup,
        )

        with self._make_client(project="project") as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "mutate_rows"
                ) as mock_gapic:
                    mock_gapic.side_effect = retryable_exception("mock")
                    with pytest.raises(MutationsExceptionGroup) as e:
                        mutation = mutations.SetCell(
                            "family", b"qualifier", b"value", timestamp_micros=123
                        )
                        entry = mutations.RowMutationEntry(b"row_key", [mutation])
                        assert mutation.is_idempotent() is True
                        table.bulk_mutate_rows([entry], operation_timeout=0.05)
                    assert len(e.value.exceptions) == 1
                    failed_exception = e.value.exceptions[0]
                    assert isinstance(failed_exception, FailedMutationEntryError)
                    assert "non-idempotent" not in str(failed_exception)
                    cause = failed_exception.__cause__
                    assert isinstance(cause, RetryExceptionGroup)
                    assert isinstance(cause.exceptions[0], retryable_exception)

    @pytest.mark.parametrize(
        "retryable_exception",
        [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable],
    )
    def test_bulk_mutate_rows_non_idempotent_retryable_errors(
        self, retryable_exception
    ):
        """Non-Idempotent mutations should never be retried"""
        from google.cloud.bigtable.data.exceptions import (
            FailedMutationEntryError,
            MutationsExceptionGroup,
        )

        with self._make_client(project="project") as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "mutate_rows"
                ) as mock_gapic:
                    mock_gapic.side_effect = lambda *a, **k: self._mock_response(
                        [retryable_exception("mock")]
                    )
                    with pytest.raises(MutationsExceptionGroup) as e:
                        mutation = mutations.SetCell(
                            "family", b"qualifier", b"value", -1
                        )
                        entry = mutations.RowMutationEntry(b"row_key", [mutation])
                        assert mutation.is_idempotent() is False
                        table.bulk_mutate_rows([entry], operation_timeout=0.2)
                    assert len(e.value.exceptions) == 1
                    failed_exception = e.value.exceptions[0]
                    assert isinstance(failed_exception, FailedMutationEntryError)
                    assert "non-idempotent" in str(failed_exception)
                    cause = failed_exception.__cause__
                    assert isinstance(cause, retryable_exception)

    @pytest.mark.parametrize(
        "non_retryable_exception",
        [
            core_exceptions.OutOfRange,
            core_exceptions.NotFound,
            core_exceptions.FailedPrecondition,
            RuntimeError,
            ValueError,
        ],
    )
    def test_bulk_mutate_rows_non_retryable_errors(self, non_retryable_exception):
        """If the request fails with a non-retryable error, mutations should not be retried"""
        from google.cloud.bigtable.data.exceptions import (
            FailedMutationEntryError,
            MutationsExceptionGroup,
        )

        with self._make_client(project="project") as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "mutate_rows"
                ) as mock_gapic:
                    mock_gapic.side_effect = non_retryable_exception("mock")
                    with pytest.raises(MutationsExceptionGroup) as e:
                        mutation = mutations.SetCell(
                            "family", b"qualifier", b"value", timestamp_micros=123
                        )
                        entry = mutations.RowMutationEntry(b"row_key", [mutation])
                        assert mutation.is_idempotent() is True
                        table.bulk_mutate_rows([entry], operation_timeout=0.2)
                    assert len(e.value.exceptions) == 1
                    failed_exception = e.value.exceptions[0]
                    assert isinstance(failed_exception, FailedMutationEntryError)
                    assert "non-idempotent" not in str(failed_exception)
                    cause = failed_exception.__cause__
                    assert isinstance(cause, non_retryable_exception)

    def test_bulk_mutate_error_index(self):
        """Test partial failure, partial success. Errors should be associated with the correct index"""
        from google.api_core.exceptions import (
            DeadlineExceeded,
            ServiceUnavailable,
            FailedPrecondition,
        )
        from google.cloud.bigtable.data.exceptions import (
            RetryExceptionGroup,
            FailedMutationEntryError,
            MutationsExceptionGroup,
        )

        with self._make_client(project="project") as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "mutate_rows"
                ) as mock_gapic:
                    mock_gapic.side_effect = [
                        self._mock_response([None, ServiceUnavailable("mock"), None]),
                        self._mock_response([DeadlineExceeded("mock")]),
                        self._mock_response([FailedPrecondition("final")]),
                    ]
                    with pytest.raises(MutationsExceptionGroup) as e:
                        mutation = mutations.SetCell(
                            "family", b"qualifier", b"value", timestamp_micros=123
                        )
                        entries = [
                            mutations.RowMutationEntry(
                                f"row_key_{i}".encode(), [mutation]
                            )
                            for i in range(3)
                        ]
                        assert mutation.is_idempotent() is True
                        table.bulk_mutate_rows(entries, operation_timeout=1000)
                    assert len(e.value.exceptions) == 1
                    failed = e.value.exceptions[0]
                    assert isinstance(failed, FailedMutationEntryError)
                    assert failed.index == 1
                    assert failed.entry == entries[1]
                    cause = failed.__cause__
                    assert isinstance(cause, RetryExceptionGroup)
                    assert len(cause.exceptions) == 3
                    assert isinstance(cause.exceptions[0], ServiceUnavailable)
                    assert isinstance(cause.exceptions[1], DeadlineExceeded)
                    assert isinstance(cause.exceptions[2], FailedPrecondition)

    def test_bulk_mutate_error_recovery(self):
        """If an error occurs, then resolves, no exception should be raised"""
        from google.api_core.exceptions import DeadlineExceeded

        with self._make_client(project="project") as client:
            table = client.get_table("instance", "table")
            with mock.patch.object(client._gapic_client, "mutate_rows") as mock_gapic:
                mock_gapic.side_effect = [
                    self._mock_response([DeadlineExceeded("mock")]),
                    self._mock_response([None]),
                ]
                mutation = mutations.SetCell(
                    "family", b"qualifier", b"value", timestamp_micros=123
                )
                entries = [
                    mutations.RowMutationEntry(f"row_key_{i}".encode(), [mutation])
                    for i in range(3)
                ]
                table.bulk_mutate_rows(entries, operation_timeout=1000)


class TestCheckAndMutateRow:
    def _make_client(self, *args, **kwargs):
        return CrossSync._Sync_Impl.TestBigtableDataClient._make_client(*args, **kwargs)

    @pytest.mark.parametrize("gapic_result", [True, False])
    def test_check_and_mutate(self, gapic_result):
        from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse

        app_profile = "app_profile_id"
        with self._make_client() as client:
            with client.get_table(
                "instance", "table", app_profile_id=app_profile
            ) as table:
                with mock.patch.object(
                    client._gapic_client, "check_and_mutate_row"
                ) as mock_gapic:
                    mock_gapic.return_value = CheckAndMutateRowResponse(
                        predicate_matched=gapic_result
                    )
                    row_key = b"row_key"
                    predicate = None
                    true_mutations = [DeleteAllFromRow()]
                    false_mutations = [DeleteAllFromRow(), DeleteAllFromRow()]
                    operation_timeout = 0.2
                    found = table.check_and_mutate_row(
                        row_key,
                        predicate,
                        true_case_mutations=true_mutations,
                        false_case_mutations=false_mutations,
                        operation_timeout=operation_timeout,
                    )
                    assert found == gapic_result
                    kwargs = mock_gapic.call_args[1]
                    request = kwargs["request"]
                    assert request.table_name == table.table_name
                    assert request.row_key == row_key
                    assert bool(request.predicate_filter) is False
                    assert request.true_mutations == [
                        m._to_pb() for m in true_mutations
                    ]
                    assert request.false_mutations == [
                        m._to_pb() for m in false_mutations
                    ]
                    assert request.app_profile_id == app_profile
                    assert kwargs["timeout"] == operation_timeout
                    assert kwargs["retry"] is None

    def test_check_and_mutate_bad_timeout(self):
        """Should raise error if operation_timeout < 0"""
        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with pytest.raises(ValueError) as e:
                    table.check_and_mutate_row(
                        b"row_key",
                        None,
                        true_case_mutations=[mock.Mock()],
                        false_case_mutations=[],
                        operation_timeout=-1,
                    )
                assert str(e.value) == "operation_timeout must be greater than 0"

    def test_check_and_mutate_single_mutations(self):
        """if single mutations are passed, they should be internally wrapped in a list"""
        from google.cloud.bigtable.data.mutations import SetCell
        from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse

        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "check_and_mutate_row"
                ) as mock_gapic:
                    mock_gapic.return_value = CheckAndMutateRowResponse(
                        predicate_matched=True
                    )
                    true_mutation = SetCell("family", b"qualifier", b"value")
                    false_mutation = SetCell("family", b"qualifier", b"value")
                    table.check_and_mutate_row(
                        b"row_key",
                        None,
                        true_case_mutations=true_mutation,
                        false_case_mutations=false_mutation,
                    )
                    kwargs = mock_gapic.call_args[1]
                    request = kwargs["request"]
                    assert request.true_mutations == [true_mutation._to_pb()]
                    assert request.false_mutations == [false_mutation._to_pb()]

    def test_check_and_mutate_predicate_object(self):
        """predicate filter should be passed to gapic request"""
        from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse
        from google.cloud.bigtable_v2.types.data import RowFilter

        mock_predicate = mock.Mock()
        predicate_pb = RowFilter({"sink": True})
        mock_predicate._to_pb.return_value = predicate_pb
        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "check_and_mutate_row"
                ) as mock_gapic:
                    mock_gapic.return_value = CheckAndMutateRowResponse(
                        predicate_matched=True
                    )
                    table.check_and_mutate_row(
                        b"row_key",
                        mock_predicate,
                        false_case_mutations=[DeleteAllFromRow()],
                    )
                    kwargs = mock_gapic.call_args[1]
                    request = kwargs["request"]
                    assert request.predicate_filter == predicate_pb
                    assert mock_predicate._to_pb.call_count == 1
                    assert kwargs["retry"] is None

    def test_check_and_mutate_mutations_parsing(self):
        """mutations objects should be converted to protos"""
        from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse
        from google.cloud.bigtable.data.mutations import DeleteAllFromFamily

        mutations = [mock.Mock() for _ in range(5)]
        for idx, mutation in enumerate(mutations):
            mutation._to_pb.return_value = DeleteAllFromFamily(f"fake {idx}")._to_pb()
        mutations.append(DeleteAllFromRow())
        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "check_and_mutate_row"
                ) as mock_gapic:
                    mock_gapic.return_value = CheckAndMutateRowResponse(
                        predicate_matched=True
                    )
                    table.check_and_mutate_row(
                        b"row_key",
                        None,
                        true_case_mutations=mutations[0:2],
                        false_case_mutations=mutations[2:],
                    )
                    kwargs = mock_gapic.call_args[1]
                    request = kwargs["request"]
                    assert request.true_mutations == [
                        DeleteAllFromFamily("fake 0")._to_pb(),
                        DeleteAllFromFamily("fake 1")._to_pb(),
                    ]
                    assert request.false_mutations == [
                        DeleteAllFromFamily("fake 2")._to_pb(),
                        DeleteAllFromFamily("fake 3")._to_pb(),
                        DeleteAllFromFamily("fake 4")._to_pb(),
                        DeleteAllFromRow()._to_pb(),
                    ]
                    assert all(
                        (mutation._to_pb.call_count == 1 for mutation in mutations[:5])
                    )


class TestReadModifyWriteRow:
    def _make_client(self, *args, **kwargs):
        return CrossSync._Sync_Impl.TestBigtableDataClient._make_client(*args, **kwargs)

    @pytest.mark.parametrize(
        "call_rules,expected_rules",
        [
            (
                AppendValueRule("f", "c", b"1"),
                [AppendValueRule("f", "c", b"1")._to_pb()],
            ),
            (
                [AppendValueRule("f", "c", b"1")],
                [AppendValueRule("f", "c", b"1")._to_pb()],
            ),
            (IncrementRule("f", "c", 1), [IncrementRule("f", "c", 1)._to_pb()]),
            (
                [AppendValueRule("f", "c", b"1"), IncrementRule("f", "c", 1)],
                [
                    AppendValueRule("f", "c", b"1")._to_pb(),
                    IncrementRule("f", "c", 1)._to_pb(),
                ],
            ),
        ],
    )
    def test_read_modify_write_call_rule_args(self, call_rules, expected_rules):
        """Test that the gapic call is called with given rules"""
        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with mock.patch.object(
                    client._gapic_client, "read_modify_write_row"
                ) as mock_gapic:
                    table.read_modify_write_row("key", call_rules)
                assert mock_gapic.call_count == 1
                found_kwargs = mock_gapic.call_args_list[0][1]
                request = found_kwargs["request"]
                assert request.rules == expected_rules
                assert found_kwargs["retry"] is None

    @pytest.mark.parametrize("rules", [[], None])
    def test_read_modify_write_no_rules(self, rules):
        with self._make_client() as client:
            with client.get_table("instance", "table") as table:
                with pytest.raises(ValueError) as e:
                    table.read_modify_write_row("key", rules=rules)
                    assert e.value.args[0] == "rules must contain at least one item"

    def test_read_modify_write_call_defaults(self):
        instance = "instance1"
        table_id = "table1"
        project = "project1"
        row_key = "row_key1"
        with self._make_client(project=project) as client:
            with client.get_table(instance, table_id) as table:
                with mock.patch.object(
                    client._gapic_client, "read_modify_write_row"
                ) as mock_gapic:
                    table.read_modify_write_row(row_key, IncrementRule("f", "q"))
                    assert mock_gapic.call_count == 1
                    kwargs = mock_gapic.call_args_list[0][1]
                    request = kwargs["request"]
                    assert (
                        request.table_name
                        == f"projects/{project}/instances/{instance}/tables/{table_id}"
                    )
                    assert bool(request.app_profile_id) is False
                    assert request.row_key == row_key.encode()
                    assert kwargs["timeout"] > 1

    def test_read_modify_write_call_overrides(self):
        row_key = b"row_key1"
        expected_timeout = 12345
        profile_id = "profile1"
        with self._make_client() as client:
            with client.get_table(
                "instance", "table_id", app_profile_id=profile_id
            ) as table:
                with mock.patch.object(
                    client._gapic_client, "read_modify_write_row"
                ) as mock_gapic:
                    table.read_modify_write_row(
                        row_key,
                        IncrementRule("f", "q"),
                        operation_timeout=expected_timeout,
                    )
                    assert mock_gapic.call_count == 1
                    kwargs = mock_gapic.call_args_list[0][1]
                    request = kwargs["request"]
                    assert request.app_profile_id == profile_id
                    assert request.row_key == row_key
                    assert kwargs["timeout"] == expected_timeout

    def test_read_modify_write_string_key(self):
        row_key = "string_row_key1"
        with self._make_client() as client:
            with client.get_table("instance", "table_id") as table:
                with mock.patch.object(
                    client._gapic_client, "read_modify_write_row"
                ) as mock_gapic:
                    table.read_modify_write_row(row_key, IncrementRule("f", "q"))
                    assert mock_gapic.call_count == 1
                    kwargs = mock_gapic.call_args_list[0][1]
                    request = kwargs["request"]
                    assert request.row_key == row_key.encode()

    def test_read_modify_write_row_building(self):
        """results from gapic call should be used to construct row"""
        from google.cloud.bigtable.data.row import Row
        from google.cloud.bigtable_v2.types import ReadModifyWriteRowResponse
        from google.cloud.bigtable_v2.types import Row as RowPB

        mock_response = ReadModifyWriteRowResponse(row=RowPB())
        with self._make_client() as client:
            with client.get_table("instance", "table_id") as table:
                with mock.patch.object(
                    client._gapic_client, "read_modify_write_row"
                ) as mock_gapic:
                    with mock.patch.object(Row, "_from_pb") as constructor_mock:
                        mock_gapic.return_value = mock_response
                        table.read_modify_write_row("key", IncrementRule("f", "q"))
                        assert constructor_mock.call_count == 1
                        constructor_mock.assert_called_once_with(mock_response.row)


class TestExecuteQuery:
    TABLE_NAME = "TABLE_NAME"
    INSTANCE_NAME = "INSTANCE_NAME"

    @pytest.fixture(scope="function")
    def client(self, *args, **kwargs):
        return CrossSync._Sync_Impl.TestBigtableDataClient._make_client(*args, **kwargs)

    @pytest.fixture(scope="function")
    def execute_query_mock(self, client):
        with mock.patch.object(
            client._gapic_client, "execute_query", CrossSync._Sync_Impl.Mock()
        ) as execute_query_mock:
            yield execute_query_mock

    @pytest.fixture(scope="function")
    def prepare_mock(self, client):
        with mock.patch.object(
            client._gapic_client, "prepare_query", CrossSync._Sync_Impl.Mock()
        ) as prepare_mock:
            prepare_mock.return_value = prepare_response(
                prepared_query=b"foo",
                metadata=metadata(column("a", str_type()), column("b", int64_type())),
            )
            yield prepare_mock

    def _make_gapic_stream(self, sample_list: list["ExecuteQueryResponse" | Exception]):
        class MockStream:
            def __init__(self, sample_list):
                self.sample_list = sample_list

            def __aiter__(self):
                return self

            def __iter__(self):
                return self

            def __next__(self):
                if not self.sample_list:
                    raise CrossSync._Sync_Impl.StopIteration
                value = self.sample_list.pop(0)
                if isinstance(value, Exception):
                    raise value
                return value

            def __anext__(self):
                return self.__next__()

        return MockStream(sample_list)

    def test_execute_query(self, client, execute_query_mock, prepare_mock):
        values = [
            *chunked_responses(2, str_val("test"), int_val(8), reset=True, token=b"r1"),
            *chunked_responses(2, str_val("test2"), int_val(9), token=b"r2"),
            *chunked_responses(2, str_val("test3"), null_val(), token=b"r3"),
        ]
        execute_query_mock.return_value = self._make_gapic_stream(values)
        result = client.execute_query(
            f"SELECT a, b FROM {self.TABLE_NAME}", self.INSTANCE_NAME
        )
        results = [r for r in result]
        assert results[0]["a"] == "test"
        assert results[0]["b"] == 8
        assert results[1]["a"] == "test2"
        assert results[1]["b"] == 9
        assert results[2]["a"] == "test3"
        assert results[2]["b"] is None
        assert execute_query_mock.call_count == 1
        assert prepare_mock.call_count == 1

    def test_execute_query_with_params(self, client, execute_query_mock, prepare_mock):
        values = [*chunked_responses(2, str_val("test2"), int_val(9), token=b"r2")]
        execute_query_mock.return_value = self._make_gapic_stream(values)
        result = client.execute_query(
            f"SELECT a, b FROM {self.TABLE_NAME} WHERE b=@b",
            self.INSTANCE_NAME,
            parameters={"b": 9},
        )
        results = [r for r in result]
        assert len(results) == 1
        assert results[0]["a"] == "test2"
        assert results[0]["b"] == 9
        assert execute_query_mock.call_count == 1
        assert prepare_mock.call_count == 1

    def test_execute_query_error_before_metadata(
        self, client, execute_query_mock, prepare_mock
    ):
        from google.api_core.exceptions import DeadlineExceeded

        values = [
            DeadlineExceeded(""),
            *chunked_responses(2, str_val("test"), int_val(8), reset=True, token=b"r1"),
            *chunked_responses(2, str_val("test2"), int_val(9), token=b"r2"),
            *chunked_responses(2, str_val("test3"), null_val(), token=b"r3"),
        ]
        execute_query_mock.return_value = self._make_gapic_stream(values)
        result = client.execute_query(
            f"SELECT a, b FROM {self.TABLE_NAME}", self.INSTANCE_NAME
        )
        results = [r for r in result]
        assert len(results) == 3
        assert execute_query_mock.call_count == 2
        assert prepare_mock.call_count == 1

    def test_execute_query_error_after_metadata(
        self, client, execute_query_mock, prepare_mock
    ):
        from google.api_core.exceptions import DeadlineExceeded

        values = [
            DeadlineExceeded(""),
            *chunked_responses(2, str_val("test"), int_val(8), reset=True, token=b"r1"),
            *chunked_responses(2, str_val("test2"), int_val(9), token=b"r2"),
            *chunked_responses(2, str_val("test3"), null_val(), token=b"r3"),
        ]
        execute_query_mock.return_value = self._make_gapic_stream(values)
        result = client.execute_query(
            f"SELECT a, b FROM {self.TABLE_NAME}", self.INSTANCE_NAME
        )
        results = [r for r in result]
        assert len(results) == 3
        assert execute_query_mock.call_count == 2
        requests = [args[0][0] for args in execute_query_mock.call_args_list]
        resume_tokens = [r.resume_token for r in requests if r.resume_token]
        assert resume_tokens == []

    def test_execute_query_with_retries(self, client, execute_query_mock, prepare_mock):
        from google.api_core.exceptions import DeadlineExceeded

        values = [
            *chunked_responses(2, str_val("test"), int_val(8), reset=True, token=b"r1"),
            DeadlineExceeded(""),
            *chunked_responses(2, str_val("test2"), int_val(9), token=b"r2"),
            DeadlineExceeded(""),
            *chunked_responses(2, str_val("test3"), null_val(), token=b"r3"),
        ]
        execute_query_mock.return_value = self._make_gapic_stream(values)
        result = client.execute_query(
            f"SELECT a, b FROM {self.TABLE_NAME}", self.INSTANCE_NAME
        )
        results = [r for r in result]
        assert results[0]["a"] == "test"
        assert results[0]["b"] == 8
        assert results[1]["a"] == "test2"
        assert results[1]["b"] == 9
        assert results[2]["a"] == "test3"
        assert results[2]["b"] is None
        assert len(results) == 3
        requests = [args[0][0] for args in execute_query_mock.call_args_list]
        resume_tokens = [r.resume_token for r in requests if r.resume_token]
        assert resume_tokens == [b"r1", b"r2"]
        assert prepare_mock.call_count == 1

    @pytest.mark.parametrize(
        "exception",
        [
            core_exceptions.DeadlineExceeded(""),
            core_exceptions.Aborted(""),
            core_exceptions.ServiceUnavailable(""),
        ],
    )
    def test_execute_query_retryable_error(
        self, client, execute_query_mock, prepare_mock, exception
    ):
        [res1, res2] = chunked_responses(
            2, str_val("test"), int_val(8), reset=True, token=b"t1"
        )
        values = [
            *chunked_responses(1, str_val("test"), int_val(8), reset=True, token=b"t1"),
            exception,
            *chunked_responses(1, str_val("tes2"), int_val(9), reset=True, token=b"t1"),
        ]
        execute_query_mock.return_value = self._make_gapic_stream(values)
        result = client.execute_query(
            f"SELECT a, b FROM {self.TABLE_NAME}", self.INSTANCE_NAME
        )
        results = [r for r in result]
        assert len(results) == 2
        assert execute_query_mock.call_count == 2
        assert prepare_mock.call_count == 1
        requests = [args[0][0] for args in execute_query_mock.call_args_list]
        resume_tokens = [r.resume_token for r in requests if r.resume_token]
        assert resume_tokens == [b"t1"]

    @pytest.mark.parametrize(
        "ExceptionType",
        [
            core_exceptions.InvalidArgument,
            core_exceptions.FailedPrecondition,
            core_exceptions.PermissionDenied,
            core_exceptions.MethodNotImplemented,
            core_exceptions.Cancelled,
            core_exceptions.AlreadyExists,
            core_exceptions.OutOfRange,
            core_exceptions.DataLoss,
            core_exceptions.Unauthenticated,
            core_exceptions.NotFound,
            core_exceptions.ResourceExhausted,
            core_exceptions.Unknown,
            core_exceptions.InternalServerError,
        ],
    )
    def test_execute_query_non_retryable(
        self, client, execute_query_mock, prepare_mock, ExceptionType
    ):
        values = [
            *chunked_responses(2, str_val("test"), int_val(8), reset=True, token=b"r1"),
            ExceptionType(""),
            *chunked_responses(2, str_val("test2"), int_val(9), token=b"r2"),
            *chunked_responses(2, str_val("test3"), null_val(), token=b"r3"),
        ]
        execute_query_mock.return_value = self._make_gapic_stream(values)
        result = client.execute_query(
            f"SELECT a, b FROM {self.TABLE_NAME}", self.INSTANCE_NAME
        )
        r = CrossSync._Sync_Impl.next(result)
        assert r["a"] == "test"
        assert r["b"] == 8
        with pytest.raises(ExceptionType):
            r = CrossSync._Sync_Impl.next(result)
        assert execute_query_mock.call_count == 1
        assert prepare_mock.call_count == 1
        requests = [args[0][0] for args in execute_query_mock.call_args_list]
        resume_tokens = [r.resume_token for r in requests if r.resume_token]
        assert resume_tokens == []

    @pytest.mark.parametrize(
        "retryable_exception",
        [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable],
    )
    def test_prepare_query_retryable(
        self, client, execute_query_mock, prepare_mock, retryable_exception
    ):
        prepare_mock.reset_mock()
        prepare_mock.side_effect = [
            retryable_exception("test"),
            prepare_response(
                b"foo",
                metadata=metadata(column("a", str_type()), column("b", int64_type())),
            ),
        ]
        values = [
            *chunked_responses(1, str_val("test"), int_val(8), reset=True, token=b"t1")
        ]
        execute_query_mock.return_value = self._make_gapic_stream(values)
        result = client.execute_query(
            f"SELECT a, b FROM {self.TABLE_NAME}", self.INSTANCE_NAME
        )
        results = [r for r in result]
        assert results[0]["a"] == "test"
        assert results[0]["b"] == 8
        assert execute_query_mock.call_count == 1
        assert prepare_mock.call_count == 2

    @pytest.mark.parametrize(
        "non_retryable_exception",
        [
            core_exceptions.InvalidArgument,
            core_exceptions.FailedPrecondition,
            core_exceptions.PermissionDenied,
            core_exceptions.MethodNotImplemented,
            core_exceptions.Cancelled,
            core_exceptions.AlreadyExists,
            core_exceptions.OutOfRange,
            core_exceptions.DataLoss,
            core_exceptions.Unauthenticated,
            core_exceptions.NotFound,
            core_exceptions.ResourceExhausted,
            core_exceptions.Unknown,
            core_exceptions.InternalServerError,
        ],
    )
    def test_prepare_query_non_retryable(
        self, client, execute_query_mock, prepare_mock, non_retryable_exception
    ):
        prepare_mock.reset_mock()
        prepare_mock.side_effect = [
            non_retryable_exception("test"),
            prepare_response(
                b"foo",
                metadata=metadata(column("a", str_type()), column("b", int64_type())),
            ),
        ]
        values = [
            *chunked_responses(1, str_val("test"), int_val(8), reset=True, token=b"t1")
        ]
        execute_query_mock.return_value = self._make_gapic_stream(values)
        with pytest.raises(non_retryable_exception):
            client.execute_query(
                f"SELECT a, b FROM {self.TABLE_NAME}", self.INSTANCE_NAME
            )