tests: Tests for sync client + fixes + client library versioning

Author: gkevinzheng

google/cloud/bigtable/admin_v2/overlay/services/bigtable_table_admin/client.py

@@ -27,10 +27,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
+import copy
 import functools
 
-from typing import Optional, Sequence, Tuple, Union
+from typing import Callable, Optional, Sequence, Tuple, Union
 from google.api_core import gapic_v1
 from google.api_core import retry as retries
 
@@ -39,16 +39,97 @@
 except AttributeError:  # pragma: NO COVER
     OptionalRetry = Union[retries.Retry, object, None]  # type: ignore
 
-from google.api_core import operation  # type: ignore
+from google.api_core import client_options as client_options_lib
+from google.auth import credentials as ga_credentials  # type: ignore
+
+from google.cloud.bigtable import __version__ as bigtable_version
 from google.cloud.bigtable.admin_v2.types import bigtable_table_admin
 
 from google.cloud.bigtable.admin_v2.services.bigtable_table_admin import (
-    BaseBigtableTableAdminClient,
+    client as base_client,
+)
+from google.cloud.bigtable.admin_v2.services.bigtable_table_admin.transports.base import (
+    BigtableTableAdminTransport
 )
 from google.cloud.bigtable.admin_v2.overlay.types import consistency, restore_table
 
 
-class BigtableTableAdminClient(BaseBigtableTableAdminClient):
+DEFAULT_CLIENT_INFO = copy.copy(base_client.DEFAULT_CLIENT_INFO)
+DEFAULT_CLIENT_INFO.client_library_version = f"{bigtable_version}-admin-overlay"
+
+
+class BigtableTableAdminClient(base_client.BaseBigtableTableAdminClient):
+    def __init__(
+        self,
+        *,
+        credentials: Optional[ga_credentials.Credentials] = None,
+        transport: Optional[
+            Union[
+                str,
+                BigtableTableAdminTransport,
+                Callable[..., BigtableTableAdminTransport],
+            ]
+        ] = None,
+        client_options: Optional[Union[client_options_lib.ClientOptions, dict]] = None,
+        client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO,
+    ) -> None:
+        """Instantiates the Bigtable table admin client.
+
+        Args:
+            credentials (Optional[google.auth.credentials.Credentials]): The
+                authorization credentials to attach to requests. These
+                credentials identify the application to the service; if none
+                are specified, the client will attempt to ascertain the
+                credentials from the environment.
+            transport (Optional[Union[str,BigtableTableAdminTransport,Callable[..., BigtableTableAdminTransport]]]):
+                The transport to use, or a Callable that constructs and returns a new transport.
+                If a Callable is given, it will be called with the same set of initialization
+                arguments as used in the BigtableTableAdminTransport constructor.
+                If set to None, a transport is chosen automatically.
+            client_options (Optional[Union[google.api_core.client_options.ClientOptions, dict]]):
+                Custom options for the client.
+
+                1. The ``api_endpoint`` property can be used to override the
+                default endpoint provided by the client when ``transport`` is
+                not explicitly provided. Only if this property is not set and
+                ``transport`` was not explicitly provided, the endpoint is
+                determined by the GOOGLE_API_USE_MTLS_ENDPOINT environment
+                variable, which have one of the following values:
+                "always" (always use the default mTLS endpoint), "never" (always
+                use the default regular endpoint) and "auto" (auto-switch to the
+                default mTLS endpoint if client certificate is present; this is
+                the default value).
+
+                2. If the GOOGLE_API_USE_CLIENT_CERTIFICATE environment variable
+                is "true", then the ``client_cert_source`` property can be used
+                to provide a client certificate for mTLS transport. If
+                not provided, the default SSL client certificate will be used if
+                present. If GOOGLE_API_USE_CLIENT_CERTIFICATE is "false" or not
+                set, no client certificate will be used.
+
+                3. The ``universe_domain`` property can be used to override the
+                default "googleapis.com" universe. Note that the ``api_endpoint``
+                property still takes precedence; and ``universe_domain`` is
+                currently not supported for mTLS.
+
+            client_info (google.api_core.gapic_v1.client_info.ClientInfo):
+                The client info used to send a user-agent string along with
+                API requests. If ``None``, then default info will be used.
+                Generally, you only need to set this if you're developing
+                your own client library.
+
+        Raises:
+            google.auth.exceptions.MutualTLSChannelError: If mutual TLS transport
+                creation failed for any reason.
+        """
+        super(BigtableTableAdminClient, self).__init__(
+            credentials=credentials,
+            transport=transport,
+            client_options=client_options,
+            client_info=client_info,
+        )
+
+
     def restore_table(
         self,
         request: Optional[Union[bigtable_table_admin.RestoreTableRequest, dict]] = None,
@@ -143,6 +224,7 @@ def sample_restore_table():
         )
         return restore_table_operation
 
+
     def wait_for_consistency(
         self,
         request: Optional[
@@ -236,6 +318,7 @@ def sample_wait_for_consistency():
         )
         return consistency.CheckConsistencyPollingFuture(api_call, default_retry=retry)
 
+
     def wait_for_replication(
         self,
         request: Optional[
@@ -334,6 +417,6 @@ def sample_wait_for_replication():
             generate_consistency_response.consistency_token
         )
 
-        return self.await_consistency(
+        return self.wait_for_consistency(
             check_consistency_request, retry=retry, timeout=timeout, metadata=metadata
         )

google/cloud/bigtable/admin_v2/overlay/types/consistency.py

@@ -17,7 +17,7 @@
 from google.api_core.future import polling
 from google.api_core import gapic_v1
 from google.api_core import retry as retries
-from google.bigtable.admin_v2.types import bigtable_table_admin
+from google.cloud.bigtable.admin_v2.types import bigtable_table_admin
 
 try:
     OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault, None]
@@ -89,6 +89,7 @@ def done(self, retry: OptionalRetry = None):
             return check_consistency_response.consistent
         except Exception as e:
             self.set_exception(e)
+            raise e
 
     def cancel(self):
         raise NotImplementedError("Cannot cancel consistency token operation")

tests/unit/admin_overlay/test_client.py

@@ -0,0 +1,219 @@
+# -*- coding: utf-8 -*-
+# Copyright 2025 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.
+
+# try/except added for compatibility with python < 3.8
+try:
+    from unittest import mock
+    from unittest.mock import AsyncMock  # pragma: NO COVER
+except ImportError:  # pragma: NO COVER
+    import mock
+
+import functools
+
+from google.api_core import gapic_v1
+from google.api_core import retry as retries
+from google.cloud.bigtable.admin_v2.services.bigtable_table_admin import transports
+from google.cloud.bigtable.admin_v2.types import bigtable_table_admin
+from google.cloud.bigtable.admin_v2.overlay.services.bigtable_table_admin.client import (
+    BigtableTableAdminClient, DEFAULT_CLIENT_INFO
+)
+from google.cloud.bigtable.admin_v2.overlay.types import consistency, restore_table
+
+import pytest
+
+
+PARENT_NAME = "my_parent"
+TABLE_NAME = "my_table"
+CONSISTENCY_TOKEN = "abcdefg"
+
+
+@pytest.mark.parametrize(
+    "transport_class,transport_name",
+    [
+        (
+            transports.BigtableTableAdminGrpcTransport,
+            "grpc",
+        ),
+        (
+            transports.BigtableTableAdminRestTransport,
+            "rest",
+        ),
+    ],
+)
+def test_bigtable_table_admin_client_client_version(
+    transport_class, transport_name
+):
+    with mock.patch.object(transport_class, "__init__") as patched:
+        patched.return_value = None
+        BigtableTableAdminClient(transport=transport_name)
+        transport_init_call = patched.call_args
+        assert transport_init_call.kwargs["client_info"] == DEFAULT_CLIENT_INFO
+
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {
+            "request": bigtable_table_admin.RestoreTableRequest(
+                parent=PARENT_NAME,
+                table_id=TABLE_NAME,
+            )
+        },
+        {
+            "request": {
+                "parent": PARENT_NAME,
+                "table_id": TABLE_NAME,
+            },
+        },
+        {
+            "request": bigtable_table_admin.RestoreTableRequest(
+                parent=PARENT_NAME,
+                table_id=TABLE_NAME,
+            ),
+            "retry": mock.Mock(spec=retries.Retry),
+            "timeout": mock.Mock(spec=retries.Retry),
+            "metadata": [("foo", "bar")]
+        },
+    ]
+)
+def test_bigtable_table_admin_client_restore_table(kwargs):
+    client = BigtableTableAdminClient()
+
+    with mock.patch.object(restore_table, "RestoreTableOperation") as future_mock:
+        with mock.patch.object(client, "_transport") as transport_mock:
+            with mock.patch.object(client, "_restore_table") as restore_table_mock:
+                operation_mock = mock.Mock()
+                restore_table_mock.return_value = operation_mock
+                client.restore_table(**kwargs)
+
+                restore_table_mock.assert_called_once_with(
+                    request=kwargs["request"],
+                    retry=kwargs.get("retry", gapic_v1.method.DEFAULT),
+                    timeout=kwargs.get("timeout", gapic_v1.method.DEFAULT),
+                    metadata=kwargs.get("metadata", ()),
+                )
+                future_mock.assert_called_once_with(transport_mock.operations_client, operation_mock)
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {
+            "request": bigtable_table_admin.CheckConsistencyRequest(
+                name=TABLE_NAME,
+                consistency_token=CONSISTENCY_TOKEN,
+            )
+        },
+        {
+            "request": {
+                "name": TABLE_NAME,
+                "consistency_token": CONSISTENCY_TOKEN,
+            },
+        },
+        {
+            "name": TABLE_NAME,
+            "consistency_token": CONSISTENCY_TOKEN,
+        },
+        {
+            "request": bigtable_table_admin.CheckConsistencyRequest(
+                name=TABLE_NAME,
+                consistency_token=CONSISTENCY_TOKEN,
+            ),
+            "retry": mock.Mock(spec=retries.Retry),
+            "timeout": mock.Mock(spec=retries.Retry),
+            "metadata": [("foo", "bar")]
+        },
+    ]
+)
+def test_bigtable_table_admin_client_wait_for_consistency(kwargs):
+    client = BigtableTableAdminClient()
+
+    future = client.wait_for_consistency(**kwargs)
+
+    assert type(future) == consistency.CheckConsistencyPollingFuture
+    assert future._default_retry == kwargs.get("retry", gapic_v1.method.DEFAULT)
+
+    check_consistency_call = future._check_consistency_call
+    assert type(check_consistency_call) == functools.partial
+
+    assert check_consistency_call.func == client.check_consistency
+    assert check_consistency_call.args == (
+        kwargs.get("request", None),
+    )
+    assert check_consistency_call.keywords == {
+        "name": kwargs.get("name", None),
+        "consistency_token": kwargs.get("consistency_token", None),
+        "timeout": kwargs.get("timeout", gapic_v1.method.DEFAULT),
+        "metadata": kwargs.get("metadata", ()),
+    }
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {
+            "request": bigtable_table_admin.GenerateConsistencyTokenRequest(
+                name=TABLE_NAME,
+            )
+        },
+        {
+            "request": {
+                "name": TABLE_NAME
+            },
+        },
+        {
+            "name": TABLE_NAME,
+        },
+        {
+            "request": bigtable_table_admin.GenerateConsistencyTokenRequest(
+                name=TABLE_NAME,
+            ),
+            "retry": mock.Mock(spec=retries.Retry),
+            "timeout": mock.Mock(spec=retries.Retry),
+            "metadata": [("foo", "bar")]
+        },
+    ]
+)
+def test_bigtable_table_admin_client_wait_for_replication(kwargs):
+    client = BigtableTableAdminClient()
+
+    with mock.patch.object(client, "generate_consistency_token") as generate_mock:
+        generate_mock.return_value = bigtable_table_admin.GenerateConsistencyTokenResponse(
+            consistency_token=CONSISTENCY_TOKEN,
+        )
+        future = client.wait_for_replication(**kwargs)
+
+        expected_check_consistency_request = bigtable_table_admin.CheckConsistencyRequest(
+            name=TABLE_NAME,
+            consistency_token=CONSISTENCY_TOKEN,
+        )
+
+        assert type(future) == consistency.CheckConsistencyPollingFuture
+        assert future._default_retry == kwargs.get("retry", gapic_v1.method.DEFAULT)
+
+        check_consistency_call = future._check_consistency_call
+        assert type(check_consistency_call) == functools.partial
+
+        assert check_consistency_call.func == client.check_consistency
+        assert check_consistency_call.args == (
+            expected_check_consistency_request,
+        )
+        assert check_consistency_call.keywords == {
+            "name": None,
+            "consistency_token": None,
+            "timeout": kwargs.get("timeout", gapic_v1.method.DEFAULT),
+            "metadata": kwargs.get("metadata", ()),
+        }