fix: registering duplicate instance (#1033)

googleapis · Nov 4, 2024 · 2bca8fb · 2bca8fb
1 parent c7e7429
commit 2bca8fb
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 1 deletion.
diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py
@@ -350,7 +350,7 @@ async def _register_instance(
             instance_name, owner.table_name, owner.app_profile_id
         )
         self._instance_owners.setdefault(instance_key, set()).add(id(owner))
-        if instance_name not in self._active_instances:
+        if instance_key not in self._active_instances:
             self._active_instances.add(instance_key)
             if self._channel_refresh_tasks:
                 # refresh tasks already running

diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py
@@ -653,6 +653,48 @@ async def test__register_instance(self):
             ]
         )
 
+    @pytest.mark.asyncio
+    async def test__register_instance_duplicate(self):
+        """
+        test double instance registration. Should be no-op
+        """
+        # set up mock client
+        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_tasks = [object()]
+        mock_channels = [mock.Mock()]
+        client_mock.transport.channels = mock_channels
+        client_mock._ping_and_warm_instances = AsyncMock()
+        table_mock = mock.Mock()
+        expected_key = (
+            "prefix/instance-1",
+            table_mock.table_name,
+            table_mock.app_profile_id,
+        )
+        # fake first registration
+        await 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])
+        # should have called ping and warm
+        assert client_mock._ping_and_warm_instances.call_count == 1
+        # next call should do nothing
+        await 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.asyncio
     @pytest.mark.parametrize(
         "insert_instances,expected_active,expected_owner_keys",