fix(tests): address high-priority review feedback

Author: hassiebp

.github/workflows/ci.yml

@@ -72,9 +72,9 @@ jobs:
 
     name: Unit tests on Python ${{ matrix.python-version }}
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
       - name: Install uv and set Python version
-        uses: astral-sh/setup-uv@v7
+        uses: astral-sh/setup-uv@cec208311dfd045dd5311c1add060b2062131d57 # v8
         with:
           version: "0.11.2"
           python-version: ${{ matrix.python-version }}

langfuse/_client/resource_manager.py

@@ -399,6 +399,8 @@ def _stop_and_join_consumer_threads(self) -> None:
         for media_upload_consumer in self._media_upload_consumers:
             media_upload_consumer.pause()
 
+        self._media_manager.signal_shutdown(count=len(self._media_upload_consumers))
+
         for media_upload_consumer in self._media_upload_consumers:
             try:
                 media_upload_consumer.join()

langfuse/_task_manager/media_manager.py

@@ -60,13 +60,14 @@ def process_next_media_upload(self) -> None:
             )
             self._queue.task_done()
 
-    def signal_shutdown(self) -> None:
-        try:
-            self._queue.put(_SHUTDOWN_SENTINEL, block=False)
-        except Full:
-            # If the queue is full, the consumer will keep draining work and
-            # observe the paused flag on the next loop iteration.
-            pass
+    def signal_shutdown(self, *, count: int = 1) -> None:
+        for _ in range(count):
+            try:
+                self._queue.put(_SHUTDOWN_SENTINEL, block=False)
+            except Full:
+                # If the queue is full, the consumer will keep draining work and
+                # observe the paused flag on the next loop iteration.
+                break
 
     def _find_and_process_media(
         self,

langfuse/_task_manager/media_upload_consumer.py

@@ -42,4 +42,3 @@ def pause(self) -> None:
             f"Thread: Pausing media upload consumer thread #{self._identifier}"
         )
         self.running = False
-        self._media_manager.signal_shutdown()

langfuse/_utils/prompt_cache.py

@@ -51,7 +51,7 @@ def run(self) -> None:
 
             if task is _SHUTDOWN_SENTINEL:
                 self._queue.task_done()
-                continue
+                break
 
             logger.debug(
                 f"PromptCacheRefreshConsumer processing task, {self._identifier}"
@@ -69,7 +69,6 @@ def run(self) -> None:
     def pause(self) -> None:
         """Pause the consumer."""
         self.running = False
-        self._queue.put(_SHUTDOWN_SENTINEL)
 
 
 class PromptCacheTaskManager(object):
@@ -134,6 +133,9 @@ def shutdown(self) -> None:
         for consumer in self._consumers:
             consumer.pause()
 
+        for _ in self._consumers:
+            self._queue.put(_SHUTDOWN_SENTINEL)
+
         for consumer in self._consumers:
             try:
                 consumer.join()

tests/conftest.py

@@ -67,7 +67,11 @@ def pytest_collection_modifyitems(items: list[pytest.Item]) -> None:
 
 
 @pytest.fixture(autouse=True)
-def reset_langfuse_state() -> Iterable[None]:
+def reset_langfuse_state(request: pytest.FixtureRequest) -> Iterable[None]:
+    if request.node.get_closest_marker("live_provider") is not None:
+        yield
+        return
+
     LangfuseResourceManager.reset()
     yield
     LangfuseResourceManager.reset()

tests/support/retry.py

@@ -52,6 +52,7 @@ def retry_until_ready(
 
             last_error = error
         else:
+            last_error = None
             if is_result_ready is None or is_result_ready(result):
                 return result
 

tests/unit/test_e2e_support.py

@@ -2,6 +2,7 @@
 
 from langfuse.api.commons.errors.not_found_error import NotFoundError
 from tests.support.api_wrapper import LangfuseAPI as SupportLangfuseAPI
+from tests.support.retry import retry_until_ready
 from tests.support.utils import get_api, wait_for_trace
 
 
@@ -138,3 +139,35 @@ class FakeClient:
     assert trace["id"] == "trace-123"
     assert len(trace["observations"]) == 3
     assert attempts["count"] == 3
+
+
+def test_retry_until_ready_clears_stale_error_after_success(monkeypatch):
+    monkeypatch.setattr("tests.support.retry.sleep", lambda _: None)
+
+    monotonic_values = iter([0.0, 0.0, 0.05, 0.06, 0.11, 0.11])
+    monkeypatch.setattr("tests.support.retry.monotonic", lambda: next(monotonic_values))
+
+    attempts = {"count": 0}
+
+    def operation():
+        attempts["count"] += 1
+
+        if attempts["count"] == 1:
+            raise NotFoundError(
+                body={
+                    "error": "LangfuseNotFoundError",
+                    "message": "Trace trace-123 not found within authorized project",
+                }
+            )
+
+        return {"id": "trace-123", "attempt": attempts["count"], "observations": []}
+
+    trace = retry_until_ready(
+        operation,
+        is_result_ready=lambda _: False,
+        timeout_seconds=0.1,
+        interval_seconds=0,
+    )
+
+    assert trace["id"] == "trace-123"
+    assert trace["attempt"] == 3

tests/unit/test_prompt.py

@@ -7,6 +7,7 @@
     DEFAULT_PROMPT_CACHE_TTL_SECONDS,
     PromptCache,
     PromptCacheItem,
+    PromptCacheTaskManager,
 )
 from langfuse.api import NotFoundError, Prompt_Chat, Prompt_Text
 from langfuse.model import ChatPromptClient, TextPromptClient
@@ -142,6 +143,42 @@ def wait_for_prompt_refresh(langfuse: Langfuse) -> None:
     langfuse._resources.prompt_cache._task_manager.wait_for_idle()
 
 
+def test_prompt_cache_task_manager_pauses_all_workers_before_broadcasting_shutdown():
+    manager = PromptCacheTaskManager(threads=0)
+    events = []
+
+    class FakeConsumer:
+        def __init__(self, identifier):
+            self._identifier = identifier
+
+        def pause(self):
+            events.append(("pause", self._identifier))
+
+        def join(self):
+            events.append(("join", self._identifier))
+
+    class FakeQueue:
+        def put(self, item):
+            events.append(("put", item))
+
+    manager._consumers = [FakeConsumer(0), FakeConsumer(1), FakeConsumer(2)]
+    manager._queue = FakeQueue()
+
+    manager.shutdown()
+
+    assert [event[0] for event in events] == [
+        "pause",
+        "pause",
+        "pause",
+        "put",
+        "put",
+        "put",
+        "join",
+        "join",
+        "join",
+    ]
+
+
 def test_get_fresh_prompt(langfuse):
     prompt_name = "test_get_fresh_prompt"
     prompt = Prompt_Text(

tests/unit/test_resource_manager.py

@@ -1,6 +1,7 @@
 """Test the LangfuseResourceManager and get_client() function."""
 
 from queue import Queue
+from types import SimpleNamespace
 from unittest.mock import Mock
 
 from langfuse import Langfuse
@@ -125,7 +126,7 @@ def test_score_ingestion_consumer_pause_wakes_blocked_thread():
     assert not consumer.is_alive()
 
 
-def test_media_upload_consumer_pause_wakes_blocked_thread():
+def test_media_upload_consumer_signal_shutdown_wakes_blocked_thread():
     media_manager = MediaManager(
         api_client=Mock(),
         httpx_client=Mock(),
@@ -135,6 +136,41 @@ def test_media_upload_consumer_pause_wakes_blocked_thread():
 
     consumer.start()
     consumer.pause()
+    media_manager.signal_shutdown()
     consumer.join(timeout=0.5)
 
     assert not consumer.is_alive()
+
+
+def test_stop_and_join_consumer_threads_broadcasts_media_shutdown_after_pausing_all():
+    events = []
+
+    class FakeConsumer:
+        def __init__(self, identifier):
+            self._identifier = identifier
+
+        def pause(self):
+            events.append(("pause", self._identifier))
+
+        def join(self):
+            events.append(("join", self._identifier))
+
+    class FakeMediaManager:
+        def signal_shutdown(self, *, count):
+            events.append(("signal_shutdown", count))
+
+    fake_resource_manager = SimpleNamespace(
+        _media_upload_consumers=[FakeConsumer(0), FakeConsumer(1)],
+        _ingestion_consumers=[],
+        _media_manager=FakeMediaManager(),
+    )
+
+    LangfuseResourceManager._stop_and_join_consumer_threads(fake_resource_manager)
+
+    assert events == [
+        ("pause", 0),
+        ("pause", 1),
+        ("signal_shutdown", 2),
+        ("join", 0),
+        ("join", 1),
+    ]