Skip to content

Commit 2bca8fb

Browse files
fix: registering duplicate instance (#1033)
1 parent c7e7429 commit 2bca8fb

File tree

2 files changed

+43
-1
lines changed

2 files changed

+43
-1
lines changed

google/cloud/bigtable/data/_async/client.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -350,7 +350,7 @@ async def _register_instance(
350350
instance_name, owner.table_name, owner.app_profile_id
351351
)
352352
self._instance_owners.setdefault(instance_key, set()).add(id(owner))
353-
if instance_name not in self._active_instances:
353+
if instance_key not in self._active_instances:
354354
self._active_instances.add(instance_key)
355355
if self._channel_refresh_tasks:
356356
# refresh tasks already running

tests/unit/data/_async/test_client.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -653,6 +653,48 @@ async def test__register_instance(self):
653653
]
654654
)
655655

656+
@pytest.mark.asyncio
657+
async def test__register_instance_duplicate(self):
658+
"""
659+
test double instance registration. Should be no-op
660+
"""
661+
# set up mock client
662+
client_mock = mock.Mock()
663+
client_mock._gapic_client.instance_path.side_effect = lambda a, b: f"prefix/{b}"
664+
active_instances = set()
665+
instance_owners = {}
666+
client_mock._active_instances = active_instances
667+
client_mock._instance_owners = instance_owners
668+
client_mock._channel_refresh_tasks = [object()]
669+
mock_channels = [mock.Mock()]
670+
client_mock.transport.channels = mock_channels
671+
client_mock._ping_and_warm_instances = AsyncMock()
672+
table_mock = mock.Mock()
673+
expected_key = (
674+
"prefix/instance-1",
675+
table_mock.table_name,
676+
table_mock.app_profile_id,
677+
)
678+
# fake first registration
679+
await self._get_target_class()._register_instance(
680+
client_mock, "instance-1", table_mock
681+
)
682+
assert len(active_instances) == 1
683+
assert expected_key == tuple(list(active_instances)[0])
684+
assert len(instance_owners) == 1
685+
assert expected_key == tuple(list(instance_owners)[0])
686+
# should have called ping and warm
687+
assert client_mock._ping_and_warm_instances.call_count == 1
688+
# next call should do nothing
689+
await self._get_target_class()._register_instance(
690+
client_mock, "instance-1", table_mock
691+
)
692+
assert len(active_instances) == 1
693+
assert expected_key == tuple(list(active_instances)[0])
694+
assert len(instance_owners) == 1
695+
assert expected_key == tuple(list(instance_owners)[0])
696+
assert client_mock._ping_and_warm_instances.call_count == 1
697+
656698
@pytest.mark.asyncio
657699
@pytest.mark.parametrize(
658700
"insert_instances,expected_active,expected_owner_keys",

0 commit comments

Comments
 (0)