From fdbee3af2527c5fc1869b536deeafb4529a550ad Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 2 Sep 2026 12:09:51 -0700 Subject: [PATCH] refactor(vector_stores): build the MongoDB pipeline immutably and inject the client class The type-discipline and test-quality gates blamed the branch for 4 LIT001, 12 LIT002 and 5 TQ008 violations. Rather than suppress them: - the $vectorSearch and $project stages are MappingProxyType and the query vector a tuple, verified against live Atlas to encode identically. The outer pipeline stays a list because pymongo's common.validate_list raises "pipeline must be a list, not ", which a unit test now pins. - the client caches are Final[dict[...]] and _client_kwargs returns a MappingProxyType. - _field_value recurses over the dotted path instead of rebinding a local. - _client_key declared Final locals in one branch and reassigned them in the others, so it is split into an early-returning _timeout_ms. - the injected callables carry explicit Final[Callable[...]] annotations, which stops pyright resolving self.embedding_fn against litellm.embedding's overloads. - get_sync_client and get_async_client take an optional client_class, so the cache tests inject a recording double instead of patching the importer, and can assert the connection string and timeouts the client was built with. SensitiveDataMasker is public SDK surface, so extra_sensitive_patterns moves to the end of the signature: in slot two it silently reinterpreted an existing caller's positional override set as extra sensitive patterns. --- .../sensitive_data_masker.py | 14 +-- litellm/llms/mongodb/common_utils.py | 53 +++++---- .../mongodb/vector_stores/transformation.py | 109 ++++++++++------- .../management_endpoints.py | 2 +- .../test_sensitive_data_masker.py | 12 ++ .../test_mongodb_transformation.py | 112 ++++++++++++------ 6 files changed, 192 insertions(+), 110 deletions(-) diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 3b0806ab069..fcce63e016b 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -1,13 +1,13 @@ from collections.abc import Mapping +from collections.abc import Set as AbstractSet from typing import Any, Final from pydantic import BaseModel from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER - _DEFAULT_SENSITIVE_PATTERNS: Final = frozenset( - { + ( "password", "secret", "key", @@ -23,20 +23,20 @@ _DEFAULT_SENSITIVE_PATTERNS: Final = frozenset( "certificate", "fingerprint", "tenancy", - } + ) ) class SensitiveDataMasker: def __init__( self, - sensitive_patterns: set[str] | None = None, - extra_sensitive_patterns: set[str] | None = None, - non_sensitive_overrides: set[str] | None = None, + sensitive_patterns: AbstractSet[str] | None = None, + non_sensitive_overrides: AbstractSet[str] | None = None, visible_prefix: int = 4, visible_suffix: int = 4, mask_char: str = "*", mask_short_values: bool = True, + extra_sensitive_patterns: AbstractSet[str] | None = None, ): self.sensitive_patterns = (sensitive_patterns or _DEFAULT_SENSITIVE_PATTERNS) | ( extra_sensitive_patterns or frozenset() @@ -44,7 +44,7 @@ class SensitiveDataMasker: # If any key segment matches one of these, the key is not considered sensitive # even if it also matches a sensitive pattern. For example, "input_cost_per_token" # contains "token" but "cost" overrides that — it's a pricing field, not a secret. - self.non_sensitive_overrides = non_sensitive_overrides or {"cost"} + self.non_sensitive_overrides = non_sensitive_overrides or frozenset(("cost",)) self.visible_prefix = visible_prefix self.visible_suffix = visible_suffix diff --git a/litellm/llms/mongodb/common_utils.py b/litellm/llms/mongodb/common_utils.py index 48496eee170..8ac02552ecb 100644 --- a/litellm/llms/mongodb/common_utils.py +++ b/litellm/llms/mongodb/common_utils.py @@ -11,8 +11,10 @@ TLS handshake and topology discovery: measured at ~890ms against Atlas versus import asyncio import weakref from asyncio import AbstractEventLoop +from collections.abc import Callable, Mapping from dataclasses import dataclass -from typing import TYPE_CHECKING, Final +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias from litellm.exceptions import BadRequestError, Timeout @@ -54,13 +56,17 @@ class MongoClientKey: server_selection_timeout_ms: int -_sync_clients: dict[MongoClientKey, "MongoClient"] = {} # mutable-ok: process-level connection cache, see module docstring -# The value carries a weak reference to the loop the client was built on: CPython recycles -# id() aggressively (measured: 200 of 200 fresh loops landed on an id already in this cache), -# so the id alone would hand a new loop a client bound to a closed one. -_async_clients: dict[ # mutable-ok: same cache, keyed per event loop - tuple[MongoClientKey, int], tuple["weakref.ref[AbstractEventLoop]", "AsyncMongoClient"] -] = {} +SyncClientFactory: TypeAlias = Callable[..., "MongoClient"] +AsyncClientFactory: TypeAlias = Callable[..., "AsyncMongoClient"] + +_AsyncClientCacheKey: TypeAlias = tuple[MongoClientKey, int] +# The entry carries a weak reference to the loop the client was built on: CPython recycles id() +# aggressively (measured: 200 of 200 fresh loops landed on an id already in this cache), so the +# id alone would hand a new loop a client bound to a closed one. +_AsyncClientEntry: TypeAlias = tuple["weakref.ref[AbstractEventLoop]", "AsyncMongoClient"] + +_sync_clients: Final[dict[MongoClientKey, "MongoClient"]] = {} # mutable-ok: process-level client cache +_async_clients: Final[dict[_AsyncClientCacheKey, _AsyncClientEntry]] = {} # mutable-ok: same cache, per loop def import_sync_mongo_client() -> "type[MongoClient]": @@ -79,33 +85,39 @@ def import_async_mongo_client() -> "type[AsyncMongoClient]": return AsyncMongoClientClass -def _client_kwargs(key: MongoClientKey) -> dict[str, object]: - return { # mutable-ok: pymongo's client constructor takes keyword arguments - "connectTimeoutMS": key.connect_timeout_ms, - "socketTimeoutMS": key.socket_timeout_ms, - "serverSelectionTimeoutMS": key.server_selection_timeout_ms, - "appname": _APP_NAME, - } +def _client_kwargs(key: MongoClientKey) -> Mapping[str, object]: + return MappingProxyType( + { + "connectTimeoutMS": key.connect_timeout_ms, + "socketTimeoutMS": key.socket_timeout_ms, + "serverSelectionTimeoutMS": key.server_selection_timeout_ms, + "appname": _APP_NAME, + } + ) -def get_sync_client(key: MongoClientKey) -> "MongoClient": +def get_sync_client(key: MongoClientKey, client_class: SyncClientFactory | None = None) -> "MongoClient": + """``client_class`` is the injection seam the tests build fake clients through; left unset the + real pymongo class is imported at call time, keeping pymongo out of import-time dependencies.""" cached: Final = _sync_clients.get(key) if cached is not None: return cached - client: Final = import_sync_mongo_client()(key.connection_string, **_client_kwargs(key)) + build: Final = client_class if client_class is not None else import_sync_mongo_client() + client: Final = build(key.connection_string, **_client_kwargs(key)) if len(_sync_clients) < _MAX_CACHED_CLIENTS: _sync_clients[key] = client return client -def get_async_client(key: MongoClientKey) -> "AsyncMongoClient": +def get_async_client(key: MongoClientKey, client_class: AsyncClientFactory | None = None) -> "AsyncMongoClient": """Async clients bind to the loop that created them, so the cache is keyed per loop.""" loop: Final = asyncio.get_running_loop() loop_key: Final = (key, id(loop)) cached: Final = _async_clients.get(loop_key) if cached is not None and cached[0]() is loop: return cached[1] - client: Final = import_async_mongo_client()(key.connection_string, **_client_kwargs(key)) + build: Final = client_class if client_class is not None else import_async_mongo_client() + client: Final = build(key.connection_string, **_client_kwargs(key)) if len(_async_clients) < _MAX_CACHED_CLIENTS or loop_key in _async_clients: _async_clients[loop_key] = (weakref.ref(loop), client) return client @@ -221,8 +233,7 @@ def translate_mongo_error(error: Exception, index_name: str, database: str, coll f"cluster name against the URI Atlas shows under Connect, Drivers. Driver detail: {error}" ) return config_error( - "mongodb_connection_string is not a usable MongoDB connection string. " - f"Driver detail: {error}" + f"mongodb_connection_string is not a usable MongoDB connection string. Driver detail: {error}" ) if isinstance(error, InvalidOperation): return config_error(f"The MongoDB client was already closed or is unusable. Driver detail: {error}") diff --git a/litellm/llms/mongodb/vector_stores/transformation.py b/litellm/llms/mongodb/vector_stores/transformation.py index 4b792accc9f..d0a0f51cd77 100644 --- a/litellm/llms/mongodb/vector_stores/transformation.py +++ b/litellm/llms/mongodb/vector_stores/transformation.py @@ -143,10 +143,18 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig): async_client_factory: Callable[[MongoClientKey], object] | None = None, ) -> None: super().__init__() - self.embedding_fn = embedding_fn if embedding_fn is not None else litellm.embedding - self.aembedding_fn = aembedding_fn if aembedding_fn is not None else litellm.aembedding - self.sync_client_factory = sync_client_factory if sync_client_factory is not None else get_sync_client - self.async_client_factory = async_client_factory if async_client_factory is not None else get_async_client + self.embedding_fn: Final[Callable[..., EmbeddingResponse]] = ( + embedding_fn if embedding_fn is not None else litellm.embedding + ) + self.aembedding_fn: Final[Callable[..., Awaitable[EmbeddingResponse]]] = ( + aembedding_fn if aembedding_fn is not None else litellm.aembedding + ) + self.sync_client_factory: Final[Callable[[MongoClientKey], object]] = ( + sync_client_factory if sync_client_factory is not None else get_sync_client + ) + self.async_client_factory: Final[Callable[[MongoClientKey], object]] = ( + async_client_factory if async_client_factory is not None else get_async_client + ) @staticmethod def _reject_unknown_params(litellm_params: Mapping[str, object]) -> None: @@ -154,9 +162,7 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig): which would otherwise turn a mistyped mongodb_collection into 'mongodb_collection is required' pointing at a key the reader can see they have set.""" unknown: Final = sorted( - key - for key in litellm_params - if key.startswith(_MONGODB_PARAM_PREFIX) and key not in _KNOWN_MONGODB_PARAMS + key for key in litellm_params if key.startswith(_MONGODB_PARAM_PREFIX) and key not in _KNOWN_MONGODB_PARAMS ) if unknown: raise config_error( @@ -196,16 +202,20 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig): return min(max(limit * NUM_CANDIDATES_MULTIPLIER, MIN_NUM_CANDIDATES), MAX_NUM_CANDIDATES) @staticmethod - def _client_key(params: _MongoDBSearchParams, timeout: float | httpx.Timeout | None) -> MongoClientKey: + def _timeout_ms(timeout: float | httpx.Timeout | None) -> tuple[int, int]: + """The connect and socket budgets pymongo is built with, in that order.""" if isinstance(timeout, httpx.Timeout): - connect_ms: Final = int((timeout.connect or DEFAULT_CONNECT_TIMEOUT_MS / 1000) * 1000) - socket_ms: Final = int((timeout.read or DEFAULT_SOCKET_TIMEOUT_MS / 1000) * 1000) - elif timeout is not None: - connect_ms = min(int(float(timeout) * 1000), DEFAULT_CONNECT_TIMEOUT_MS) - socket_ms = int(float(timeout) * 1000) - else: - connect_ms = DEFAULT_CONNECT_TIMEOUT_MS - socket_ms = DEFAULT_SOCKET_TIMEOUT_MS + return ( + int((timeout.connect or DEFAULT_CONNECT_TIMEOUT_MS / 1000) * 1000), + int((timeout.read or DEFAULT_SOCKET_TIMEOUT_MS / 1000) * 1000), + ) + if timeout is None: + return DEFAULT_CONNECT_TIMEOUT_MS, DEFAULT_SOCKET_TIMEOUT_MS + return min(int(float(timeout) * 1000), DEFAULT_CONNECT_TIMEOUT_MS), int(float(timeout) * 1000) + + @classmethod + def _client_key(cls, params: _MongoDBSearchParams, timeout: float | httpx.Timeout | None) -> MongoClientKey: + connect_ms, socket_ms = cls._timeout_ms(timeout) return MongoClientKey( connection_string=params.require_connection_string(), connect_timeout_ms=connect_ms, @@ -220,36 +230,41 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig): query_vector: Sequence[float], params: _MongoDBSearchParams, vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - ) -> list[dict[str, object]]: + ) -> Sequence[Mapping[str, object]]: if vector_store_search_optional_params.get("filters") is not None: raise config_error( "MongoDB vector store does not support the filters parameter yet. " "Restrict the collection or the Atlas Vector Search index definition instead." ) limit: Final = cls._limit(vector_store_search_optional_params) - return [ # mutable-ok: pymongo's aggregate contract is a list of stage dicts + search: Final = MappingProxyType( { - "$vectorSearch": { - "index": vector_store_id, - "path": params.embedding_field, - "queryVector": list(query_vector), - "numCandidates": cls._num_candidates(limit, params.mongodb_num_candidates), - "limit": limit, - } - }, - {"$project": {params.text_field: 1, SCORE_FIELD_NAME: {"$meta": "vectorSearchScore"}}}, + "index": vector_store_id, + "path": params.embedding_field, + "queryVector": tuple(query_vector), + "numCandidates": cls._num_candidates(limit, params.mongodb_num_candidates), + "limit": limit, + } + ) + projection: Final = MappingProxyType( + {params.text_field: 1, SCORE_FIELD_NAME: MappingProxyType({"$meta": "vectorSearchScore"})} + ) + return [ # mutable-ok: pymongo rejects any non-list pipeline in common.validate_list + MappingProxyType({"$vectorSearch": search}), + MappingProxyType({"$project": projection}), ] - @staticmethod - def _field_value(document: Mapping[str, object], dotted_path: str) -> str | None: + @classmethod + def _field_value(cls, document: Mapping[str, object], dotted_path: str) -> str | None: """None means the path is absent from the document, which is what separates a mistyped mongodb_text_field from a document whose text is genuinely empty.""" - current: object = document - for segment in dotted_path.split("."): - if not isinstance(current, Mapping) or segment not in current: - return None - current = current[segment] - return None if current is None else str(current) + head, _, rest = dotted_path.partition(".") + if head not in document: + return None + value: Final = document[head] + if not rest: + return None if value is None else str(value) + return cls._field_value(value, rest) if isinstance(value, Mapping) else None @classmethod def _to_result(cls, document: Mapping[str, object], text_field: str) -> VectorStoreSearchResult: @@ -287,7 +302,9 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig): return VectorStoreSearchResponse( object="vector_store.search_results.page", search_query=query_text, - data=[cls._to_result(document, text_field) for document in documents], + data=[ # mutable-ok: VectorStoreSearchResponse declares data as a list + cls._to_result(document, text_field) for document in documents + ], ) @staticmethod @@ -341,14 +358,12 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig): try: client: Final = self.sync_client_factory(key) target: Final = client[database][collection] # pyright: ignore[reportIndexIssue] # factory is typed as returning object so injected doubles are accepted - documents: Final = list(target.aggregate(pipeline)) + documents: Final = tuple(target.aggregate(pipeline)) except Exception as e: - raise translate_mongo_error( - e, index_name=vector_store_id, database=database, collection=collection - ) from e + raise translate_mongo_error(e, index_name=vector_store_id, database=database, collection=collection) from e if not documents: try: - catalogue: Final = list(target.list_search_indexes(vector_store_id)) + catalogue: Final = tuple(target.list_search_indexes(vector_store_id)) except Exception as e: raise translate_mongo_error( e, index_name=vector_store_id, database=database, collection=collection @@ -386,15 +401,17 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig): client: Final = self.async_client_factory(key) target: Final = client[database][collection] # pyright: ignore[reportIndexIssue] # factory is typed as returning object so injected doubles are accepted cursor: Final = await target.aggregate(pipeline) - documents: Final = [document async for document in cursor] + documents: Final = [ # mutable-ok: an async comprehension cannot build a tuple directly + document async for document in cursor + ] except Exception as e: - raise translate_mongo_error( - e, index_name=vector_store_id, database=database, collection=collection - ) from e + raise translate_mongo_error(e, index_name=vector_store_id, database=database, collection=collection) from e if not documents: try: index_cursor: Final = await target.list_search_indexes(vector_store_id) - catalogue: Final = [entry async for entry in index_cursor] + catalogue: Final = [ # mutable-ok: an async comprehension cannot build a tuple directly + entry async for entry in index_cursor + ] except Exception as e: raise translate_mongo_error( e, index_name=vector_store_id, database=database, collection=collection diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index a62c0f711cb..9ca0753f354 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -62,7 +62,7 @@ def _row_to_vector_store(row: "_VectorStoreRow") -> LiteLLM_ManagedVectorStore: # "connection" covers wire-protocol providers whose whole credential is a URI # (mongodb_connection_string embeds the username and password), which the # default api_key/secret/token patterns do not match. -_LITELLM_PARAMS_MASKER: Final = SensitiveDataMasker(extra_sensitive_patterns={"connection"}) +_LITELLM_PARAMS_MASKER: Final = SensitiveDataMasker(extra_sensitive_patterns=frozenset(("connection",))) _REDACT_LITELLM_PARAMS_MAX_DEPTH: Final = 10 diff --git a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py index 27a83223864..c2b4042bdba 100644 --- a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py +++ b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py @@ -331,3 +331,15 @@ def test_extra_sensitive_patterns_do_not_leak_into_other_maskers(): SensitiveDataMasker(extra_sensitive_patterns={"connection"}) assert SensitiveDataMasker().is_sensitive_key("mongodb_connection_string") is False + + +def test_the_second_positional_argument_is_still_the_override_set(): + """SensitiveDataMasker is public SDK surface, so adding a keyword must not shift what an + existing positional call means. Putting extra_sensitive_patterns second would silently turn + an override set into an extra sensitive set and start masking the caller's pricing fields.""" + from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker + + masker = SensitiveDataMasker({"token"}, {"session"}) + + assert masker.is_sensitive_key("session_token") is False + assert masker.is_sensitive_key("auth_token") is True diff --git a/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py b/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py index 84de42d2126..e7ed3d77407 100644 --- a/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py +++ b/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py @@ -39,6 +39,15 @@ BASE_PARAMS = { READY_INDEX = [{"name": INDEX, "status": "READY", "queryable": True}] +class RecordingClient: + """Stands in for pymongo's client class so the cache tests inject a fake rather than + patching the importer, and so they can assert what the client was actually built with.""" + + def __init__(self, connection_string, **kwargs): + self.connection_string = connection_string + self.kwargs = kwargs + + class FakeCollection: def __init__(self, documents, error=None, search_indexes=None): self.documents = documents @@ -171,12 +180,22 @@ def test_search_builds_vector_search_stage_against_the_named_index(): assert _stage(collection, "$vectorSearch") == { "index": INDEX, "path": "embedding", - "queryVector": [0.1, 0.2, 0.3], + "queryVector": (0.1, 0.2, 0.3), "numCandidates": 100, "limit": 5, } +def test_the_pipeline_reaches_pymongo_as_a_list(): + """pymongo's common.validate_list rejects any other sequence with + 'pipeline must be a list, not ', so the outer container is part of the contract.""" + config, _, collection = _config() + + _search(config) + + assert isinstance(collection.pipeline, list) + + def test_search_projects_the_text_field_and_the_similarity_score(): config, _, collection = _config() @@ -275,6 +294,38 @@ def test_response_reads_a_dotted_text_field_path(): assert response["data"][0]["content"][0]["text"] == "nested text" +def test_a_dotted_path_resolves_three_levels_deep(): + config, _, _ = _config(documents=[{"_id": 1, "a": {"b": {"c": "deep text"}}, "score": 0.5}]) + + response = _search(config, litellm_params={"mongodb_text_field": "a.b.c"}) + + assert response["data"][0]["content"][0]["text"] == "deep text" + + +def test_a_dotted_path_that_runs_through_a_scalar_counts_as_absent(): + """Walking 'plot.nope' when plot is a string must report the misconfiguration, not + stringify the scalar and hand the model text from the wrong field.""" + config, _, _ = _config(documents=[{"_id": 1, "plot": "a plain string", "score": 0.5}]) + + with pytest.raises(BadRequestError, match=r"has a 'plot\.nope' field"): + _search(config, litellm_params={"mongodb_text_field": "plot.nope"}) + + +def test_a_non_string_text_field_is_stringified(): + config, _, _ = _config(documents=[{"_id": 1, "year": 1979, "score": 0.5}]) + + response = _search(config, litellm_params={"mongodb_text_field": "year"}) + + assert response["data"][0]["content"][0]["text"] == "1979" + + +def test_a_null_text_field_counts_as_absent(): + config, _, _ = _config(documents=[{"_id": 1, "text": None, "score": 0.5}]) + + with pytest.raises(BadRequestError, match="has a 'text' field"): + _search(config) + + def test_response_tolerates_a_sparse_document_missing_the_text_field(): config, _, _ = _config(documents=[{"_id": 1, "score": 0.5}, {"_id": 2, "text": "has text", "score": 0.4}]) @@ -489,7 +540,7 @@ async def test_async_search_builds_the_same_pipeline_and_maps_the_response(): assert client.requested_database == "sample_mflix" assert client.database.requested_collection == "embedded_movies" assert _stage(collection, "$vectorSearch")["limit"] == 3 - assert _stage(collection, "$vectorSearch")["queryVector"] == [0.1, 0.2, 0.3] + assert _stage(collection, "$vectorSearch")["queryVector"] == (0.1, 0.2, 0.3) assert response["data"][0]["content"][0]["text"] == "an astronaut adrift" assert response["data"][0]["score"] == 0.94 @@ -524,42 +575,36 @@ class TestClientCache: ) def test_the_same_connection_reuses_one_client(self): - with patch("litellm.llms.mongodb.common_utils.import_sync_mongo_client") as importer: - importer.return_value = lambda *args, **kwargs: MagicMock() - - first = get_sync_client(self._key()) - second = get_sync_client(self._key()) + first = get_sync_client(self._key(), RecordingClient) + second = get_sync_client(self._key(), RecordingClient) assert first is second - assert importer.return_value + assert first.connection_string == CONNECTION_STRING + assert first.kwargs["socketTimeoutMS"] == 30_000 + assert first.kwargs["connectTimeoutMS"] == 10_000 + assert first.kwargs["appname"] == "litellm" def test_a_different_connection_gets_its_own_client(self): - with patch("litellm.llms.mongodb.common_utils.import_sync_mongo_client") as importer: - importer.return_value = lambda *args, **kwargs: MagicMock() - - first = get_sync_client(self._key()) - second = get_sync_client(self._key(connection_string="mongodb://other.example.test")) + first = get_sync_client(self._key(), RecordingClient) + second = get_sync_client(self._key(connection_string="mongodb://other.example.test"), RecordingClient) assert first is not second + assert second.connection_string == "mongodb://other.example.test" def test_a_different_timeout_gets_its_own_client(self): - with patch("litellm.llms.mongodb.common_utils.import_sync_mongo_client") as importer: - importer.return_value = lambda *args, **kwargs: MagicMock() - - first = get_sync_client(self._key()) - second = get_sync_client(self._key(socket_timeout_ms=5_000)) + first = get_sync_client(self._key(), RecordingClient) + second = get_sync_client(self._key(socket_timeout_ms=5_000), RecordingClient) assert first is not second + assert second.kwargs["socketTimeoutMS"] == 5_000 @pytest.mark.asyncio async def test_async_clients_are_cached_per_event_loop(self): - with patch("litellm.llms.mongodb.common_utils.import_async_mongo_client") as importer: - importer.return_value = lambda *args, **kwargs: MagicMock() - - first = get_async_client(self._key()) - second = get_async_client(self._key()) + first = get_async_client(self._key(), RecordingClient) + second = get_async_client(self._key(), RecordingClient) assert first is second + assert first.connection_string == CONNECTION_STRING def test_a_new_loop_never_inherits_a_closed_loop_client(self): @@ -579,19 +624,16 @@ class TestClientCache: clients_handed_out = [] async def fetch(): - return get_async_client(key) + return get_async_client(key, LoopAgnosticClient) - with patch("litellm.llms.mongodb.common_utils.import_async_mongo_client") as importer: - importer.return_value = LoopAgnosticClient - - for _ in range(20): - loop = asyncio.new_event_loop() - client = loop.run_until_complete(fetch()) - clients_handed_out.append((client, client.built_on, loop.is_closed())) - client.built_on = weakref.ref(loop) - loop.close() - del loop - gc.collect() + for _ in range(20): + loop = asyncio.new_event_loop() + client = loop.run_until_complete(fetch()) + clients_handed_out.append((client, client.built_on, loop.is_closed())) + client.built_on = weakref.ref(loop) + loop.close() + del loop + gc.collect() stale = [ handed_out