linting + name changes in tests + added test for timeout

Author: gkevinzheng

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

@@ -42,17 +42,18 @@
 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 (
     client as base_client,
 )
 from google.cloud.bigtable.admin_v2.services.bigtable_table_admin.transports.base import (
-    BigtableTableAdminTransport
+    BigtableTableAdminTransport,
 )
 from google.cloud.bigtable.admin_v2.overlay.types import consistency, restore_table
 
+from google.cloud.bigtable.gapic_version import __version__ as bigtable_version
+
 
 DEFAULT_CLIENT_INFO = copy.copy(base_client.DEFAULT_CLIENT_INFO)
 DEFAULT_CLIENT_INFO.client_library_version = f"{bigtable_version}-admin-overlay"
@@ -129,7 +130,6 @@ def __init__(
             client_info=client_info,
         )
 
-
     def restore_table(
         self,
         request: Optional[Union[bigtable_table_admin.RestoreTableRequest, dict]] = None,
@@ -224,7 +224,6 @@ def sample_restore_table():
         )
         return restore_table_operation
 
-
     def wait_for_consistency(
         self,
         request: Optional[
@@ -318,7 +317,6 @@ def sample_wait_for_consistency():
         )
         return consistency.CheckConsistencyPollingFuture(api_call, default_retry=retry)
 
-
     def wait_for_replication(
         self,
         request: Optional[

tests/unit/admin_overlay/test_client.py

@@ -27,10 +27,13 @@
 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
+    BigtableTableAdminClient,
+    DEFAULT_CLIENT_INFO,
 )
 from google.cloud.bigtable.admin_v2.overlay.types import consistency, restore_table
 
+from google.cloud.bigtable import __version__ as bigtable_version
+
 import pytest
 
 
@@ -52,15 +55,17 @@
         ),
     ],
 )
-def test_bigtable_table_admin_client_client_version(
-    transport_class, transport_name
-):
+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
 
+    assert (
+        DEFAULT_CLIENT_INFO.client_library_version
+        == f"{bigtable_version}-admin-overlay"
+    )
 
 
 @pytest.mark.parametrize(
@@ -85,9 +90,9 @@ def test_bigtable_table_admin_client_client_version(
             ),
             "retry": mock.Mock(spec=retries.Retry),
             "timeout": mock.Mock(spec=retries.Retry),
-            "metadata": [("foo", "bar")]
+            "metadata": [("foo", "bar")],
         },
-    ]
+    ],
 )
 def test_bigtable_table_admin_client_restore_table(kwargs):
     client = BigtableTableAdminClient()
@@ -105,7 +110,9 @@ def test_bigtable_table_admin_client_restore_table(kwargs):
                     timeout=kwargs.get("timeout", gapic_v1.method.DEFAULT),
                     metadata=kwargs.get("metadata", ()),
                 )
-                future_mock.assert_called_once_with(transport_mock.operations_client, operation_mock)
+                future_mock.assert_called_once_with(
+                    transport_mock.operations_client, operation_mock
+                )
 
 
 @pytest.mark.parametrize(
@@ -134,9 +141,9 @@ def test_bigtable_table_admin_client_restore_table(kwargs):
             ),
             "retry": mock.Mock(spec=retries.Retry),
             "timeout": mock.Mock(spec=retries.Retry),
-            "metadata": [("foo", "bar")]
+            "metadata": [("foo", "bar")],
         },
-    ]
+    ],
 )
 def test_bigtable_table_admin_client_wait_for_consistency(kwargs):
     client = BigtableTableAdminClient()
@@ -150,9 +157,7 @@ def test_bigtable_table_admin_client_wait_for_consistency(kwargs):
     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.args == (kwargs.get("request", None),)
     assert check_consistency_call.keywords == {
         "name": kwargs.get("name", None),
         "consistency_token": kwargs.get("consistency_token", None),
@@ -170,9 +175,7 @@ def test_bigtable_table_admin_client_wait_for_consistency(kwargs):
             )
         },
         {
-            "request": {
-                "name": TABLE_NAME
-            },
+            "request": {"name": TABLE_NAME},
         },
         {
             "name": TABLE_NAME,
@@ -183,22 +186,26 @@ def test_bigtable_table_admin_client_wait_for_consistency(kwargs):
             ),
             "retry": mock.Mock(spec=retries.Retry),
             "timeout": mock.Mock(spec=retries.Retry),
-            "metadata": [("foo", "bar")]
+            "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,
+        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,
+        expected_check_consistency_request = (
+            bigtable_table_admin.CheckConsistencyRequest(
+                name=TABLE_NAME,
+                consistency_token=CONSISTENCY_TOKEN,
+            )
         )
 
         assert type(future) == consistency.CheckConsistencyPollingFuture
@@ -208,9 +215,7 @@ def test_bigtable_table_admin_client_wait_for_replication(kwargs):
         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.args == (expected_check_consistency_request,)
         assert check_consistency_call.keywords == {
             "name": None,
             "consistency_token": None,

tests/unit/admin_overlay/test_consistency.py

@@ -36,6 +36,7 @@
     consistent=False
 )
 
+
 def mock_check_consistency_callable(max_poll_count=1):
     # Return False max_poll_count - 1 times, then True, for a total of
     # max_poll_count calls.
@@ -49,7 +50,7 @@ def test_check_consistency_future_cancel():
     future = consistency.CheckConsistencyPollingFuture(check_consistency_call)
     with pytest.raises(NotImplementedError):
         future.cancel()
-    
+
     with pytest.raises(NotImplementedError):
         future.cancelled()
 
@@ -61,31 +62,37 @@ def test_check_consistency_future_result():
     is_consistent = future.result()
 
     assert is_consistent == True
-    check_consistency_call.assert_has_calls([mock.call(retry=future._default_retry)] * times)
+    check_consistency_call.assert_has_calls(
+        [mock.call(retry=future._default_retry)] * times
+    )
 
     # Check that calling result again doesn't produce more calls.
     is_consistent = future.result()
 
     assert is_consistent == True
-    check_consistency_call.assert_has_calls([mock.call(retry=future._default_retry)] * times)
+    check_consistency_call.assert_has_calls(
+        [mock.call(retry=future._default_retry)] * times
+    )
 
 
 def test_check_consistency_future_result_default_retry():
     times = 5
     check_consistency_call = mock_check_consistency_callable(times)
     retry = mock.Mock(spec=retries.Retry)
-    future = consistency.CheckConsistencyPollingFuture(check_consistency_call, default_retry=retry)
+    future = consistency.CheckConsistencyPollingFuture(
+        check_consistency_call, default_retry=retry
+    )
     is_consistent = future.result()
 
     assert is_consistent == True
     check_consistency_call.assert_has_calls([mock.call(retry=retry)] * times)
 
 
 def test_check_consistency_future_result_different_retry():
-    check_consistency_call = mock.Mock(spec=["__call__"],
-                                       side_effect=exceptions.DeadlineExceeded("Deadline Exceeded.")
+    check_consistency_call = mock.Mock(
+        spec=["__call__"], side_effect=exceptions.DeadlineExceeded("Deadline Exceeded.")
     )
     future = consistency.CheckConsistencyPollingFuture(check_consistency_call)
 
     with pytest.raises(exceptions.DeadlineExceeded):
-        future.result()
+        future.result()