From 2bca8fb220eeb1906fc6a3cf1f879f3d41fbbff8 Mon Sep 17 00:00:00 2001
From: Daniel Sanche <sanche@google.com>
Date: Mon, 4 Nov 2024 14:02:25 -0800
Subject: [PATCH] fix: registering duplicate instance (#1033)

---
 google/cloud/bigtable/data/_async/client.py |  2 +-
 tests/unit/data/_async/test_client.py       | 42 +++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py
index 82a874918..b48921623 100644
--- 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
index 6c49ca0da..1c1c14cd3 100644
--- 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",