stabilize e2e readbacks in ci

Author: hassiebp

.github/workflows/ci.yml

@@ -199,7 +199,7 @@ jobs:
       - name: Run the end-to-end tests
         run: |
           python --version
-          uv run --frozen pytest -s -v --log-cli-level=INFO tests/e2e
+          uv run --frozen pytest -n auto --dist loadfile -s -v --log-cli-level=INFO tests/e2e
 
   all-tests-passed:
     # This allows us to have a branch protection rule for tests and deploys with matrix

tests/e2e/test_core_sdk.py

@@ -152,7 +152,7 @@ def test_create_session_score():
     sleep(2)
 
     # Retrieve and verify
-    score = langfuse.api.scores.get_by_id(score_id)
+    score = get_api().scores.get_by_id(score_id)
 
     # find the score by name (server may transform the id format)
     assert score is not None
@@ -1841,26 +1841,26 @@ def test_get_observations():
 def test_get_trace_not_found():
     # Attempt to fetch a non-existent trace using the API
     with pytest.raises(Exception):
-        get_api().trace.get(create_uuid())
+        get_api(retry=False).trace.get(create_uuid())
 
 
 def test_get_observation_not_found():
     # Attempt to fetch a non-existent observation using the API
     with pytest.raises(Exception):
-        get_api().legacy.observations_v1.get(create_uuid())
+        get_api(retry=False).legacy.observations_v1.get(create_uuid())
 
 
 def test_get_traces_empty():
     # Fetch traces with a filter that should return no results
-    response = get_api().trace.list(name=create_uuid())
+    response = get_api(retry=False).trace.list(name=create_uuid())
 
     assert len(response.data) == 0
     assert response.meta.total_items == 0
 
 
 def test_get_observations_empty():
     # Fetch observations with a filter that should return no results
-    response = get_api().legacy.observations_v1.get_many(name=create_uuid())
+    response = get_api(retry=False).legacy.observations_v1.get_many(name=create_uuid())
 
     assert len(response.data) == 0
     assert response.meta.total_items == 0

tests/support/api_wrapper.py

@@ -1,8 +1,10 @@
 import os
-from time import sleep
 
 import httpx
 
+from langfuse.api.commons.errors.not_found_error import NotFoundError
+from tests.support.retry import is_not_found_payload, retry_until_ready
+
 
 class LangfuseAPI:
     def __init__(self, username=None, password=None, base_url=None):
@@ -11,28 +13,32 @@ def __init__(self, username=None, password=None, base_url=None):
         self.auth = (username, password)
         self.BASE_URL = base_url if base_url else os.environ["LANGFUSE_BASE_URL"]
 
+    def _get_json(self, url, params=None):
+        def _request():
+            response = httpx.get(url, params=params, auth=self.auth)
+            payload = response.json()
+
+            if response.status_code == 404 and is_not_found_payload(payload):
+                raise NotFoundError(body=payload, headers=dict(response.headers))
+
+            return payload
+
+        return retry_until_ready(_request)
+
     def get_observation(self, observation_id):
-        sleep(1)
         url = f"{self.BASE_URL}/api/public/observations/{observation_id}"
-        response = httpx.get(url, auth=self.auth)
-        return response.json()
+        return self._get_json(url)
 
     def get_scores(self, page=None, limit=None, user_id=None, name=None):
-        sleep(1)
         params = {"page": page, "limit": limit, "userId": user_id, "name": name}
         url = f"{self.BASE_URL}/api/public/scores"
-        response = httpx.get(url, params=params, auth=self.auth)
-        return response.json()
+        return self._get_json(url, params=params)
 
     def get_traces(self, page=None, limit=None, user_id=None, name=None):
-        sleep(1)
         params = {"page": page, "limit": limit, "userId": user_id, "name": name}
         url = f"{self.BASE_URL}/api/public/traces"
-        response = httpx.get(url, params=params, auth=self.auth)
-        return response.json()
+        return self._get_json(url, params=params)
 
     def get_trace(self, trace_id):
-        sleep(1)
         url = f"{self.BASE_URL}/api/public/traces/{trace_id}"
-        response = httpx.get(url, auth=self.auth)
-        return response.json()
+        return self._get_json(url)

tests/support/retry.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import os
+from time import monotonic, sleep
+from typing import Callable, TypeVar
+
+from langfuse.api.commons.errors.not_found_error import NotFoundError
+from langfuse.api.core.api_error import ApiError
+
+T = TypeVar("T")
+
+DEFAULT_RETRY_TIMEOUT_SECONDS = float(
+    os.environ.get("LANGFUSE_E2E_READ_TIMEOUT_SECONDS", "12")
+)
+DEFAULT_RETRY_INTERVAL_SECONDS = float(
+    os.environ.get("LANGFUSE_E2E_READ_INTERVAL_SECONDS", "0.25")
+)
+
+
+def is_eventual_consistency_error(error: Exception) -> bool:
+    if isinstance(error, NotFoundError):
+        return True
+
+    if not isinstance(error, ApiError):
+        return False
+
+    body = error.body
+    return isinstance(body, dict) and body.get("error") == "LangfuseNotFoundError"
+
+
+def is_not_found_payload(payload: object) -> bool:
+    return isinstance(payload, dict) and payload.get("error") == "LangfuseNotFoundError"
+
+
+def retry_until_ready(
+    operation: Callable[[], T],
+    *,
+    is_retryable_error: Callable[[Exception], bool] = is_eventual_consistency_error,
+    is_result_ready: Callable[[T], bool] | None = None,
+    timeout_seconds: float = DEFAULT_RETRY_TIMEOUT_SECONDS,
+    interval_seconds: float = DEFAULT_RETRY_INTERVAL_SECONDS,
+) -> T:
+    deadline = monotonic() + timeout_seconds
+    last_error: Exception | None = None
+
+    while True:
+        try:
+            result = operation()
+        except Exception as error:
+            if not is_retryable_error(error) or monotonic() >= deadline:
+                raise
+
+            last_error = error
+        else:
+            if is_result_ready is None or is_result_ready(result):
+                return result
+
+            if monotonic() >= deadline:
+                return result
+
+        sleep(interval_seconds)
+
+        if monotonic() >= deadline and last_error is not None:
+            raise last_error

tests/support/utils.py

@@ -1,23 +1,72 @@
 import base64
 import os
-from time import sleep
+from typing import Any
 from uuid import uuid4
 
 from langfuse.api import LangfuseAPI
+from tests.support.retry import retry_until_ready
+
+READ_METHOD_NAMES = {"get", "get_by_id", "get_many", "get_run", "list"}
+PAGINATION_ARGUMENTS = {"limit", "page"}
+
+
+def _has_filters(kwargs: dict[str, Any]) -> bool:
+    return any(
+        key not in PAGINATION_ARGUMENTS and value is not None
+        for key, value in kwargs.items()
+    )
+
+
+class _RetryingApiProxy:
+    def __init__(self, target: Any):
+        self._target = target
+
+    def __getattr__(self, name: str) -> Any:
+        attr = getattr(self._target, name)
+
+        if callable(attr):
+            if name not in READ_METHOD_NAMES:
+                return attr
+
+            def _call(*args: Any, **kwargs: Any) -> Any:
+                return retry_until_ready(
+                    lambda: attr(*args, **kwargs),
+                    is_result_ready=_result_ready(name, kwargs),
+                )
+
+            return _call
+
+        if isinstance(attr, (str, bytes, int, float, bool, list, dict, tuple, set)):
+            return attr
+
+        if attr is None:
+            return None
+
+        return _RetryingApiProxy(attr)
+
+
+def _result_ready(method_name: str, kwargs: dict[str, Any]):
+    if method_name not in {"get_many", "list"} or not _has_filters(kwargs):
+        return None
+
+    def _has_data(result: Any) -> bool:
+        data = getattr(result, "data", None)
+        return data is None or len(data) > 0
+
+    return _has_data
 
 
 def create_uuid():
     return str(uuid4())
 
 
-def get_api():
-    sleep(2)
-
-    return LangfuseAPI(
+def get_api(*, retry: bool = True):
+    client = LangfuseAPI(
         username=os.environ.get("LANGFUSE_PUBLIC_KEY"),
         password=os.environ.get("LANGFUSE_SECRET_KEY"),
         base_url=os.environ.get("LANGFUSE_BASE_URL"),
     )
+    return _RetryingApiProxy(client) if retry else client
 
 
 def encode_file_to_base64(image_path) -> str:

tests/unit/test_e2e_support.py

@@ -0,0 +1,116 @@
+from types import SimpleNamespace
+
+from langfuse.api.commons.errors.not_found_error import NotFoundError
+from tests.support.api_wrapper import LangfuseAPI as SupportLangfuseAPI
+from tests.support.utils import get_api
+
+
+def test_get_api_retries_not_found(monkeypatch):
+    monkeypatch.setattr("tests.support.retry.sleep", lambda _: None)
+
+    attempts = {"count": 0}
+
+    class FakeTraceService:
+        def get(self, trace_id):
+            attempts["count"] += 1
+
+            if attempts["count"] < 3:
+                raise NotFoundError(
+                    body={
+                        "error": "LangfuseNotFoundError",
+                        "message": f"Trace {trace_id} not found within authorized project",
+                    }
+                )
+
+            return {"id": trace_id}
+
+    class FakeClient:
+        trace = FakeTraceService()
+
+    monkeypatch.setattr("tests.support.utils.LangfuseAPI", lambda **_: FakeClient())
+
+    trace = get_api().trace.get("trace-123")
+
+    assert trace == {"id": "trace-123"}
+    assert attempts["count"] == 3
+
+
+def test_get_api_retries_filtered_lists(monkeypatch):
+    monkeypatch.setattr("tests.support.retry.sleep", lambda _: None)
+
+    attempts = {"count": 0}
+
+    class FakeTraceService:
+        def list(self, **kwargs):
+            attempts["count"] += 1
+
+            if attempts["count"] < 3:
+                return SimpleNamespace(data=[])
+
+            return SimpleNamespace(data=[kwargs["name"]])
+
+    class FakeClient:
+        trace = FakeTraceService()
+
+    monkeypatch.setattr("tests.support.utils.LangfuseAPI", lambda **_: FakeClient())
+
+    response = get_api().trace.list(name="ready-trace")
+
+    assert response.data == ["ready-trace"]
+    assert attempts["count"] == 3
+
+
+def test_get_api_retry_can_be_disabled(monkeypatch):
+    attempts = {"count": 0}
+
+    class FakeTraceService:
+        def list(self, **kwargs):
+            attempts["count"] += 1
+            return SimpleNamespace(data=[])
+
+    class FakeClient:
+        trace = FakeTraceService()
+
+    monkeypatch.setattr("tests.support.utils.LangfuseAPI", lambda **_: FakeClient())
+
+    response = get_api(retry=False).trace.list(name="missing-trace")
+
+    assert response.data == []
+    assert attempts["count"] == 1
+
+
+def test_raw_api_wrapper_retries_not_found_payload(monkeypatch):
+    monkeypatch.setattr("tests.support.retry.sleep", lambda _: None)
+
+    attempts = {"count": 0}
+
+    class FakeResponse:
+        def __init__(self, status_code, payload):
+            self.status_code = status_code
+            self._payload = payload
+            self.headers = {}
+
+        def json(self):
+            return self._payload
+
+    def fake_get(*args, **kwargs):
+        attempts["count"] += 1
+
+        if attempts["count"] < 3:
+            return FakeResponse(
+                404,
+                {
+                    "error": "LangfuseNotFoundError",
+                    "message": "Trace trace-123 not found within authorized project",
+                },
+            )
+
+        return FakeResponse(200, {"id": "trace-123", "observations": []})
+
+    monkeypatch.setattr("tests.support.api_wrapper.httpx.get", fake_get)
+
+    api = SupportLangfuseAPI(username="user", password="pass", base_url="http://test")
+    trace = api.get_trace("trace-123")
+
+    assert trace["id"] == "trace-123"
+    assert attempts["count"] == 3