Merge pull request #39895 from BerriAI/litellm_deflake_20260905

test: deflake redis semantic cache sys.modules leak, LangSmith init loop patch, wall-clock stagger assertion, and zombie grandchild check in the fake prisma cli
2026-09-14 23:21:35 +00:00 · 2026-09-12 15:20:26 -07:00 · 2026-09-12 15:20:26 -07:00 · d72ae3b9a2
commit d72ae3b9a2
parent fc053dae3e 441f863abb
4 changed files with 64 additions and 134 deletions
--- a/tests/test_litellm/caching/test_redis_semantic_cache.py
+++ b/tests/test_litellm/caching/test_redis_semantic_cache.py
@ -1,3 +1,5 @@
+from collections.abc import Iterator
+from contextlib import contextmanager
 from importlib import import_module
 import sys
 from unittest.mock import AsyncMock, MagicMock, patch
@ -5,18 +7,19 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import pytest


+@contextmanager
+def _fake_redisvl_modules(semantic_cache_mock: MagicMock, custom_vectorizer_mock: MagicMock) -> Iterator[None]:
+    with pytest.MonkeyPatch.context() as mp:
+        mp.setitem(sys.modules, "redisvl.extensions.llmcache", MagicMock(SemanticCache=semantic_cache_mock))
+        mp.setitem(sys.modules, "redisvl.utils.vectorize", MagicMock(CustomTextVectorizer=custom_vectorizer_mock))
+        yield
+

 # Tests for RedisSemanticCache
 def test_redis_semantic_cache_initialization(monkeypatch):
    # Mock the redisvl import
    semantic_cache_mock = MagicMock()
-    with patch.dict(
-        "sys.modules",
-        {
-            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
-            "redisvl.utils.vectorize": MagicMock(CustomTextVectorizer=MagicMock()),
-        },
-    ):
+    with _fake_redisvl_modules(semantic_cache_mock, MagicMock()):
        from litellm.caching.redis_semantic_cache import RedisSemanticCache

        # Set environment variables
@ -44,15 +47,7 @@ def test_redis_semantic_cache_get_cache(monkeypatch):
    semantic_cache_mock = MagicMock()
    custom_vectorizer_mock = MagicMock()

-    with patch.dict(
-        "sys.modules",
-        {
-            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
-            "redisvl.utils.vectorize": MagicMock(
-                CustomTextVectorizer=custom_vectorizer_mock
-            ),
-        },
-    ):
+    with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
        from litellm.caching.redis_semantic_cache import RedisSemanticCache

        # Set environment variables
@ -110,15 +105,7 @@ def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch):
    semantic_cache_mock = MagicMock()
    custom_vectorizer_mock = MagicMock()

-    with patch.dict(
-        "sys.modules",
-        {
-            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
-            "redisvl.utils.vectorize": MagicMock(
-                CustomTextVectorizer=custom_vectorizer_mock
-            ),
-        },
-    ):
+    with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
        from litellm.caching.redis_semantic_cache import RedisSemanticCache

        monkeypatch.setenv("REDIS_HOST", "localhost")
@ -162,15 +149,7 @@ def test_redis_semantic_cache_set_cache_stores_cache_key_filter(monkeypatch):
    semantic_cache_mock = MagicMock()
    custom_vectorizer_mock = MagicMock()

-    with patch.dict(
-        "sys.modules",
-        {
-            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
-            "redisvl.utils.vectorize": MagicMock(
-                CustomTextVectorizer=custom_vectorizer_mock
-            ),
-        },
-    ):
+    with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
        from litellm.caching.redis_semantic_cache import RedisSemanticCache

        monkeypatch.setenv("REDIS_HOST", "localhost")
@ -210,15 +189,7 @@ def test_redis_semantic_cache_uses_isolated_index_for_old_schema(monkeypatch):
    )
    custom_vectorizer_mock = MagicMock()

-    with patch.dict(
-        "sys.modules",
-        {
-            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
-            "redisvl.utils.vectorize": MagicMock(
-                CustomTextVectorizer=custom_vectorizer_mock
-            ),
-        },
-    ):
+    with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
        from litellm.caching.redis_semantic_cache import RedisSemanticCache

        monkeypatch.setenv("REDIS_HOST", "localhost")
@ -252,15 +223,7 @@ def test_redis_semantic_cache_overwrites_stale_isolated_index(monkeypatch):
    )
    custom_vectorizer_mock = MagicMock()

-    with patch.dict(
-        "sys.modules",
-        {
-            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
-            "redisvl.utils.vectorize": MagicMock(
-                CustomTextVectorizer=custom_vectorizer_mock
-            ),
-        },
-    ):
+    with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
        from litellm.caching.redis_semantic_cache import RedisSemanticCache

        monkeypatch.setenv("REDIS_HOST", "localhost")
@ -292,15 +255,7 @@ def test_redis_semantic_cache_reraises_unexpected_isolated_index_error(monkeypat
    )
    custom_vectorizer_mock = MagicMock()

-    with patch.dict(
-        "sys.modules",
-        {
-            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
-            "redisvl.utils.vectorize": MagicMock(
-                CustomTextVectorizer=custom_vectorizer_mock
-            ),
-        },
-    ):
+    with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
        from litellm.caching.redis_semantic_cache import RedisSemanticCache

        monkeypatch.setenv("REDIS_HOST", "localhost")
@ -369,15 +324,15 @@ def test_redis_semantic_cache_builds_filter_expression(monkeypatch):
        def __eq__(self, value):
            return (self.field_name, value)

-    with patch.dict("sys.modules", {"redisvl.query.filter": MagicMock(Tag=FakeTag)}):
-        from litellm.caching.redis_semantic_cache import RedisSemanticCache
+    monkeypatch.setitem(sys.modules, "redisvl.query.filter", MagicMock(Tag=FakeTag))
+    from litellm.caching.redis_semantic_cache import RedisSemanticCache

-        redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
+    redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)

-        assert redis_semantic_cache._get_cache_key_filter_expression("test_key") == (
-            RedisSemanticCache.CACHE_KEY_FIELD_NAME,
-            "test_key",
-        )
+    assert redis_semantic_cache._get_cache_key_filter_expression("test_key") == (
+        RedisSemanticCache.CACHE_KEY_FIELD_NAME,
+        "test_key",
+    )


@pytest.mark.asyncio
@ -386,15 +341,7 @@ async def test_redis_semantic_cache_async_get_cache(monkeypatch):
    semantic_cache_mock = MagicMock()
    custom_vectorizer_mock = MagicMock()

-    with patch.dict(
-        "sys.modules",
-        {
-            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
-            "redisvl.utils.vectorize": MagicMock(
-                CustomTextVectorizer=custom_vectorizer_mock
-            ),
-        },
-    ):
+    with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
        from litellm.caching.redis_semantic_cache import RedisSemanticCache

        # Set environment variables
@ -449,15 +396,7 @@ async def test_redis_semantic_cache_async_get_cache_rejects_unscoped_hit(monkeyp
    semantic_cache_mock = MagicMock()
    custom_vectorizer_mock = MagicMock()

-    with patch.dict(
-        "sys.modules",
-        {
-            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
-            "redisvl.utils.vectorize": MagicMock(
-                CustomTextVectorizer=custom_vectorizer_mock
-            ),
-        },
-    ):
+    with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
        from litellm.caching.redis_semantic_cache import RedisSemanticCache

        monkeypatch.setenv("REDIS_HOST", "localhost")
@ -499,15 +438,7 @@ async def test_redis_semantic_cache_async_set_cache_stores_cache_key_filter(
    semantic_cache_mock = MagicMock()
    custom_vectorizer_mock = MagicMock()

-    with patch.dict(
-        "sys.modules",
-        {
-            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
-            "redisvl.utils.vectorize": MagicMock(
-                CustomTextVectorizer=custom_vectorizer_mock
-            ),
-        },
-    ):
+    with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
        from litellm.caching.redis_semantic_cache import RedisSemanticCache

        monkeypatch.setenv("REDIS_HOST", "localhost")
@ -1255,15 +1186,7 @@ def test_redis_init_defers_redisvl_construction(monkeypatch):
    semantic_cache_mock = MagicMock()
    custom_vectorizer_mock = MagicMock()

-    with patch.dict(
-        "sys.modules",
-        {
-            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
-            "redisvl.utils.vectorize": MagicMock(
-                CustomTextVectorizer=custom_vectorizer_mock
-            ),
-        },
-    ):
+    with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
        from litellm.caching.redis_semantic_cache import RedisSemanticCache

        monkeypatch.setenv("REDIS_HOST", "localhost")
@ -1291,15 +1214,7 @@ def test_redis_failed_llmcache_build_is_not_memoized(monkeypatch):
    )
    custom_vectorizer_mock = MagicMock()

-    with patch.dict(
-        "sys.modules",
-        {
-            "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
-            "redisvl.utils.vectorize": MagicMock(
-                CustomTextVectorizer=custom_vectorizer_mock
-            ),
-        },
-    ):
+    with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
        from litellm.caching.redis_semantic_cache import RedisSemanticCache

        monkeypatch.setenv("REDIS_HOST", "localhost")
--- a/tests/test_litellm/integrations/test_langsmith_init.py
+++ b/tests/test_litellm/integrations/test_langsmith_init.py
@ -1,5 +1,6 @@
+import asyncio
 import os
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch

 import pytest

@ -154,24 +155,22 @@ class TestLangsmithLoggerInit:
        assert logger._start_periodic_flush_task() is None
        mock_get_running_loop.assert_called_once()

-    @patch("asyncio.get_running_loop")
-    def test_langsmith_init_starts_periodic_flush_with_running_loop(
-        self, mock_get_running_loop
-    ):
+    @pytest.mark.asyncio
+    async def test_langsmith_init_starts_periodic_flush_with_running_loop(self):
        """Test that init schedules periodic flush when a running loop exists."""
-        mock_loop = MagicMock()
-        mock_task = MagicMock()
-        mock_loop.create_task.return_value = mock_task
-        mock_get_running_loop.return_value = mock_loop
-
        logger = LangsmithLogger(
-            langsmith_api_key="test-key", langsmith_project="test-project"
+            langsmith_api_key="test-key", langsmith_project="test-project", flush_interval=0.01
        )
+        batch_sent = asyncio.Event()
+        logger.async_send_batch = AsyncMock(side_effect=batch_sent.set)
+        logger.log_queue.append({"id": "run-id"})

-        assert logger._flush_task == mock_task
-        mock_loop.create_task.assert_called_once()
-        scheduled_coro = mock_loop.create_task.call_args.args[0]
-        scheduled_coro.close()
+        flush_task = logger._flush_task
+        assert isinstance(flush_task, asyncio.Task)
+        await asyncio.wait_for(batch_sent.wait(), timeout=5)
+        flush_task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await flush_task

    @pytest.mark.asyncio
    async def test_async_log_success_event_lazily_starts_periodic_flush(self):
--- a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py
+++ b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py
@ -64,6 +64,10 @@ def _stagger(scheduler: AsyncIOScheduler, identity: str = "pod-a:1", **overrides
    return apply_scheduled_job_stagger(scheduler=scheduler, settings=_settings(**overrides), identity=identity)


+def _trigger_of(scheduler: AsyncIOScheduler, job_id: str):
+    return next(job.trigger for job in scheduler.get_jobs() if job.id == job_id)
+
+
 def _fire_times(trigger, start: datetime, steps: int) -> tuple[datetime, ...]:
    """The fire times APScheduler would produce, each computed from the one before it"""
    return tuple(
@ -133,22 +137,24 @@ def test_default_cron_is_staggered_and_keeps_its_offset_on_every_later_fire():
    applied = _stagger(scheduler)
    assert applied[PTU_ROLLUP_JOB_ID] > 0

-    trigger = next(job.trigger for job in scheduler.get_jobs() if job.id == PTU_ROLLUP_JOB_ID)
-    fires = _fire_times(trigger, datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc), 3)
+    start = datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)
+    fires = _fire_times(_trigger_of(scheduler, PTU_ROLLUP_JOB_ID), start, 3)

    expected = timedelta(minutes=15) + timedelta(seconds=applied[PTU_ROLLUP_JOB_ID])
    assert [fire - fire.replace(hour=0, minute=0, second=0, microsecond=0) for fire in fires] == [expected] * 3


-async def test_explicit_offset_overrides_the_derived_one_and_zero_pins_a_job():
+def test_explicit_offset_overrides_the_derived_one_and_zero_pins_a_job():
    scheduler = _with_jobs(_scheduler())
    applied = _stagger(scheduler, offsets={"periodic_reload_job": 0, PTU_ROLLUP_JOB_ID: 7})
-    unstaggered = _next_run_times(_with_jobs(_scheduler()))
-    staggered = _next_run_times(scheduler)

    assert applied["periodic_reload_job"] == 0
    assert applied[PTU_ROLLUP_JOB_ID] == 7
-    assert staggered[PTU_ROLLUP_JOB_ID] - unstaggered[PTU_ROLLUP_JOB_ID] == timedelta(seconds=7)
+
+    start = datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)
+    staggered = _trigger_of(scheduler, PTU_ROLLUP_JOB_ID)
+    unstaggered = _trigger_of(_with_jobs(_scheduler()), PTU_ROLLUP_JOB_ID)
+    assert _fire_times(staggered, start, 1)[0] - _fire_times(unstaggered, start, 1)[0] == timedelta(seconds=7)


 async def test_disabling_the_stagger_leaves_every_schedule_untouched():
--- a/tests/test_litellm/proxy/db/conftest.py
+++ b/tests/test_litellm/proxy/db/conftest.py
@ -35,6 +35,14 @@ DB_ENV_KEYS = (
 _db_env_snapshot_key = pytest.StashKey[dict[str, Optional[str]]]()


+def _is_zombie(pid: int) -> bool:
+    try:
+        stat: Final = Path(f"/proc/{pid}/stat").read_text()
+    except OSError:
+        return False
+    return stat.rpartition(")")[2].split()[0] == "Z"
+
+
 def _db_env_snapshot() -> dict[str, Optional[str]]:
    return {key: os.environ.get(key) for key in DB_ENV_KEYS}

@ -136,6 +144,8 @@ class FakePrismaCli:
                os.kill(pid, 0)
            except ProcessLookupError:
                return True
+            if _is_zombie(pid):
+                return True
            time.sleep(0.05)
        return False