speed up unit test suite

Author: hassiebp

.github/workflows/ci.yml

@@ -89,7 +89,7 @@ jobs:
       - name: Run the automated tests
         run: |
           python --version
-          uv run --frozen pytest -n auto --dist loadfile -s -v --log-cli-level=INFO tests/unit
+          uv run --frozen pytest -n auto --dist worksteal -s -v --log-cli-level=INFO tests/unit
 
   e2e-tests:
     runs-on: ubuntu-latest

langfuse/_utils/prompt_cache.py

@@ -3,7 +3,7 @@
 import atexit
 import os
 from datetime import datetime
-from queue import Empty, Queue
+from queue import Queue
 from threading import Thread
 from typing import Callable, Dict, List, Optional, Set
 
@@ -18,6 +18,7 @@
 )
 
 DEFAULT_PROMPT_CACHE_REFRESH_WORKERS = 1
+_SHUTDOWN_SENTINEL = object()
 
 
 class PromptCacheItem:
@@ -46,26 +47,29 @@ def __init__(self, queue: Queue, identifier: int):
 
     def run(self) -> None:
         while self.running:
+            task = self._queue.get()
+
+            if task is _SHUTDOWN_SENTINEL:
+                self._queue.task_done()
+                continue
+
+            logger.debug(
+                f"PromptCacheRefreshConsumer processing task, {self._identifier}"
+            )
             try:
-                task = self._queue.get(timeout=1)
-                logger.debug(
-                    f"PromptCacheRefreshConsumer processing task, {self._identifier}"
+                task()
+            # Task failed, but we still consider it processed
+            except Exception as e:
+                logger.warning(
+                    f"PromptCacheRefreshConsumer encountered an error, cache was not refreshed: {self._identifier}, {e}"
                 )
-                try:
-                    task()
-                # Task failed, but we still consider it processed
-                except Exception as e:
-                    logger.warning(
-                        f"PromptCacheRefreshConsumer encountered an error, cache was not refreshed: {self._identifier}, {e}"
-                    )
 
-                self._queue.task_done()
-            except Empty:
-                pass
+            self._queue.task_done()
 
     def pause(self) -> None:
         """Pause the consumer."""
         self.running = False
+        self._queue.put(_SHUTDOWN_SENTINEL)
 
 
 class PromptCacheTaskManager(object):
@@ -99,6 +103,9 @@ def add_task(self, key: str, task: Callable[[], None]) -> None:
     def active_tasks(self) -> int:
         return len(self._processing_keys)
 
+    def wait_for_idle(self) -> None:
+        self._queue.join()
+
     def _wrap_task(self, key: str, task: Callable[[], None]) -> Callable[[], None]:
         def wrapped() -> None:
             logger.debug(f"Refreshing prompt cache for key: {key}")

tests/unit/test_otel.py

@@ -54,10 +54,17 @@ class TestOTelBase:
     @pytest.fixture(scope="function", autouse=True)
     def cleanup_otel(self):
         """Reset OpenTelemetry state between tests."""
-        original_provider = trace_api.get_tracer_provider()
+        from opentelemetry.util._once import Once
+
+        trace_api._TRACER_PROVIDER = None
+        trace_api._PROXY_TRACER_PROVIDER = trace_api.ProxyTracerProvider()
+        trace_api._TRACER_PROVIDER_SET_ONCE = Once()
+
         yield
-        trace_api.set_tracer_provider(original_provider)
         LangfuseResourceManager.reset()
+        trace_api._TRACER_PROVIDER = None
+        trace_api._PROXY_TRACER_PROVIDER = trace_api.ProxyTracerProvider()
+        trace_api._TRACER_PROVIDER_SET_ONCE = Once()
 
     @pytest.fixture
     def memory_exporter(self):
@@ -97,7 +104,7 @@ def mock_init(self, **kwargs):
                 self,
                 span_exporter=memory_exporter,
                 max_export_batch_size=512,
-                schedule_delay_millis=5000,
+                schedule_delay_millis=1,
             )
 
         monkeypatch.setattr(
@@ -1870,7 +1877,7 @@ def update_random_metadata(thread_id):
             update = random.choice(updates)
 
             # Sleep a tiny bit to simulate work and increase chances of thread interleaving
-            time.sleep(random.uniform(0.001, 0.01))
+            time.sleep(random.uniform(0.0005, 0.001))
 
             # Apply the update to current_metadata (in a real system, this would update OTEL span)
             with metadata_lock:
@@ -2001,7 +2008,7 @@ def mock_processor_init(self, **kwargs):
                 self,
                 span_exporter=exporter,
                 max_export_batch_size=512,
-                schedule_delay_millis=5000,
+                schedule_delay_millis=1,
             )
 
         monkeypatch.setattr(
@@ -2118,7 +2125,7 @@ def create_spans_project1():
                     metadata={"project": "project1", "index": i},
                 )
                 # Small sleep to ensure overlap with other thread
-                time.sleep(0.01)
+                time.sleep(0.001)
                 span.end()
 
         def create_spans_project2():
@@ -2128,7 +2135,7 @@ def create_spans_project2():
                     metadata={"project": "project2", "index": i},
                 )
                 # Small sleep to ensure overlap with other thread
-                time.sleep(0.01)
+                time.sleep(0.001)
                 span.end()
 
         # Start threads
@@ -2378,7 +2385,7 @@ def mock_processor_init(self, **kwargs):
                 self,
                 span_exporter=exporter,
                 max_export_batch_size=512,
-                schedule_delay_millis=5000,
+                schedule_delay_millis=1,
             )
 
         monkeypatch.setattr(
@@ -2757,7 +2764,7 @@ async def async_task(parent_span, task_id):
             child_span = parent_span.start_observation(name=f"async-task-{task_id}")
 
             # Simulate async work
-            await asyncio.sleep(0.1)
+            await asyncio.sleep(0.01)
 
             # Update span with results
             child_span.update(
@@ -2948,7 +2955,7 @@ async def test_span_metadata_updates_in_async_context(
 
             # Define async tasks that update different parts of metadata
             async def update_temperature():
-                await asyncio.sleep(0.1)  # Simulate some async work
+                await asyncio.sleep(0.01)  # Simulate some async work
                 main_span.update(
                     metadata={
                         "llm_config": {
@@ -2960,7 +2967,7 @@ async def update_temperature():
                 )
 
             async def update_model():
-                await asyncio.sleep(0.05)  # Simulate some async work
+                await asyncio.sleep(0.005)  # Simulate some async work
                 main_span.update(
                     metadata={
                         "llm_config": {
@@ -2970,7 +2977,7 @@ async def update_model():
                 )
 
             async def add_context_length():
-                await asyncio.sleep(0.15)  # Simulate some async work
+                await asyncio.sleep(0.015)  # Simulate some async work
                 main_span.update(
                     metadata={
                         "llm_config": {
@@ -2982,7 +2989,7 @@ async def add_context_length():
                 )
 
             async def update_user_id():
-                await asyncio.sleep(0.08)  # Simulate some async work
+                await asyncio.sleep(0.008)  # Simulate some async work
                 main_span.update(
                     metadata={
                         "request_info": {
@@ -3047,7 +3054,7 @@ def test_metrics_and_timing(self, langfuse_client, memory_exporter):
         span = langfuse_client.start_observation(name="timing-test-span")
 
         # Add a small delay
-        time.sleep(0.1)
+        time.sleep(0.01)
 
         # End the span
         span.end()
@@ -3089,10 +3096,10 @@ def test_metrics_and_timing(self, langfuse_client, memory_exporter):
         ) / 1_000_000_000
         assert span_duration_seconds > 0, "Span duration should be positive"
 
-        # Since we slept for 0.1 seconds, the span duration should be at least 0.05 seconds
+        # Since we slept for 0.01 seconds, the span duration should be at least 0.005 seconds
         # but we'll be generous with the upper bound due to potential system delays
-        assert span_duration_seconds >= 0.05, (
-            f"Span duration ({span_duration_seconds}s) should be at least 0.05s"
+        assert span_duration_seconds >= 0.005, (
+            f"Span duration ({span_duration_seconds}s) should be at least 0.005s"
         )
 
 
@@ -3349,6 +3356,7 @@ def langfuse_client(self, monkeypatch):
             public_key="test-public-key",
             secret_key="test-secret-key",
             base_url="http://test-host",
+            tracing_enabled=False,
         )
 
         return client

tests/unit/test_prompt.py

@@ -1,4 +1,3 @@
-from time import sleep
 from unittest.mock import Mock, patch
 
 import pytest
@@ -139,6 +138,10 @@ def langfuse():
     return langfuse_instance
 
 
+def wait_for_prompt_refresh(langfuse: Langfuse) -> None:
+    langfuse._resources.prompt_cache._task_manager.wait_for_idle()
+
+
 def test_get_fresh_prompt(langfuse):
     prompt_name = "test_get_fresh_prompt"
     prompt = Prompt_Text(
@@ -376,10 +379,7 @@ def test_get_fresh_prompt_when_expired_cache_custom_ttl(mock_time, langfuse: Lan
 
     result_call_3 = langfuse.get_prompt(prompt_name)
 
-    while True:
-        if langfuse._resources.prompt_cache._task_manager.active_tasks() == 0:
-            break
-        sleep(0.1)
+    wait_for_prompt_refresh(langfuse)
 
     assert mock_server_call.call_count == 2
     assert result_call_3 == prompt_client
@@ -483,10 +483,7 @@ def test_get_stale_prompt_when_expired_cache_default_ttl(mock_time, langfuse: La
     langfuse.get_prompt(prompt_name)
     langfuse.get_prompt(prompt_name)
 
-    while True:
-        if langfuse._resources.prompt_cache._task_manager.active_tasks() == 0:
-            break
-        sleep(0.1)
+    wait_for_prompt_refresh(langfuse)
 
     assert mock_server_call.call_count == 2
 
@@ -527,10 +524,7 @@ def test_get_fresh_prompt_when_expired_cache_default_ttl(mock_time, langfuse: La
     mock_time.return_value = DEFAULT_PROMPT_CACHE_TTL_SECONDS + 1
 
     result_call_3 = langfuse.get_prompt(prompt_name)
-    while True:
-        if langfuse._resources.prompt_cache._task_manager.active_tasks() == 0:
-            break
-        sleep(0.1)
+    wait_for_prompt_refresh(langfuse)
 
     assert mock_server_call.call_count == 2
     assert result_call_3 == prompt_client
@@ -563,10 +557,7 @@ def test_get_expired_prompt_when_failing_fetch(mock_time, langfuse: Langfuse):
     mock_server_call.side_effect = Exception("Server error")
 
     result_call_2 = langfuse.get_prompt(prompt_name, max_retries=1)
-    while True:
-        if langfuse._resources.prompt_cache._task_manager.active_tasks() == 0:
-            break
-        sleep(0.1)
+    wait_for_prompt_refresh(langfuse)
 
     assert mock_server_call.call_count == 3
     assert result_call_2 == prompt_client
@@ -619,10 +610,7 @@ def raise_not_found(*_args: object, **_kwargs: object) -> None:
     )
     assert stale_result == prompt_client
 
-    while True:
-        if langfuse._resources.prompt_cache._task_manager.active_tasks() == 0:
-            break
-        sleep(0.1)
+    wait_for_prompt_refresh(langfuse)
 
     assert langfuse._resources.prompt_cache.get(cache_key) is None
 

tests/unit/test_prompt_atexit.py

@@ -20,13 +20,15 @@ def test_prompts_atexit():
 print("Adding prompt cache", PromptCache)
 prompt_cache = PromptCache(max_prompt_refresh_workers=10)
 
-# example task that takes 2 seconds but we will force it to exit earlier
-def wait_2_sec():
-    time.sleep(2)
+# example task that stays in flight briefly while the process exits
+def wait_briefly():
+    time.sleep(0.1)
 
 # 8 times
 for i in range(8):
-    prompt_cache.add_refresh_prompt_task(f"key_wait_2_sec_i_{i}", lambda: wait_2_sec())
+    prompt_cache.add_refresh_prompt_task(
+        f"key_wait_briefly_i_{i}", lambda: wait_briefly()
+    )
 """
 
     process = subprocess.Popen(
@@ -74,12 +76,14 @@ async def main():
     print("Adding prompt cache", PromptCache)
     prompt_cache = PromptCache(max_prompt_refresh_workers=10)
 
-    # example task that takes 2 seconds but we will force it to exit earlier
-    def wait_2_sec():
-        time.sleep(2)
+    # example task that stays in flight briefly while the process exits
+    def wait_briefly():
+        time.sleep(0.1)
 
     async def add_new_prompt_refresh(i: int):
-        prompt_cache.add_refresh_prompt_task(f"key_wait_2_sec_i_{i}", lambda: wait_2_sec())
+        prompt_cache.add_refresh_prompt_task(
+            f"key_wait_briefly_i_{i}", lambda: wait_briefly()
+        )
     
     # 8 times
     tasks = [add_new_prompt_refresh(i) for i in range(8)]