diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index df990c43530..9884e9d9bc0 100644 --- 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") diff --git a/tests/test_litellm/integrations/test_langsmith_init.py b/tests/test_litellm/integrations/test_langsmith_init.py index 025aa86466c..0bc9e279fbf 100644 --- 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): diff --git a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py index 54e0aa74a25..3eabfc5c840 100644 --- 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(): diff --git a/tests/test_litellm/proxy/db/conftest.py b/tests/test_litellm/proxy/db/conftest.py index d3226b0ec50..bcb7794a20d 100644 --- 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