diff --git a/pymilvus/milvus_client/async_milvus_client.py b/pymilvus/milvus_client/async_milvus_client.py index b748d0c62..fcf34c489 100644 --- a/pymilvus/milvus_client/async_milvus_client.py +++ b/pymilvus/milvus_client/async_milvus_client.py @@ -865,14 +865,25 @@ async def refresh_load( timeout: Optional[float] = None, **kwargs, ): + kwargs.pop("_refresh", None) conn = await self._get_connection() - return await conn.refresh_load( - collection_name, - partition_names, - timeout=timeout, - context=self._generate_call_context(**kwargs), - **kwargs, - ) + if partition_names: + await conn.load_partitions( + collection_name, + partition_names, + timeout=timeout, + _refresh=True, + context=self._generate_call_context(**kwargs), + **kwargs, + ) + else: + await conn.load_collection( + collection_name, + timeout=timeout, + _refresh=True, + context=self._generate_call_context(**kwargs), + **kwargs, + ) async def get_server_version( self, timeout: Optional[float] = None, detail: bool = False, **kwargs diff --git a/tests/unit/test_async_milvus_client_ops.py b/tests/unit/test_async_milvus_client_ops.py index 0873acda8..34d5ba42c 100644 --- a/tests/unit/test_async_milvus_client_ops.py +++ b/tests/unit/test_async_milvus_client_ops.py @@ -132,7 +132,6 @@ async def test_release_partitions_str_converts(self): ("add_collection_function", ("col", MagicMock()), {}, "add_collection_function"), ("alter_collection_function", ("col", "fn", MagicMock()), {}, "alter_collection_function"), # Server ops - ("refresh_load", ("col",), {}, "refresh_load"), ("run_analyzer", ("hello world",), {}, "run_analyzer"), ("update_replicate_configuration", (), {"clusters": []}, "update_replicate_configuration"), ("get_replicate_configuration", (), {}, "get_replicate_configuration"), @@ -181,6 +180,30 @@ async def test_compact_target_size_over_signed_int64_mb_rejected_before_rpc(self await client.compact("col", target_size=1 << 63) handler.compact.assert_not_called() + @pytest.mark.asyncio + async def test_refresh_load_refreshes_collection(self): + client, handler = _make_client() + + await client.refresh_load("col", timeout=3) + + handler.load_collection.assert_awaited_once_with( + "col", timeout=3, _refresh=True, context=ANY + ) + handler.load_partitions.assert_not_awaited() + handler.refresh_load.assert_not_awaited() + + @pytest.mark.asyncio + async def test_refresh_load_refreshes_partitions(self): + client, handler = _make_client() + + await client.refresh_load("col", ["part1"], timeout=3) + + handler.load_partitions.assert_awaited_once_with( + "col", ["part1"], timeout=3, _refresh=True, context=ANY + ) + handler.load_collection.assert_not_awaited() + handler.refresh_load.assert_not_awaited() + class TestAsyncClientAliasAndServerOps: @pytest.mark.asyncio