fix(tests): stabilize remaining CI failures

Author: hassiebp

tests/e2e/test_datasets.py

@@ -3,7 +3,7 @@
 
 from langfuse import Langfuse
 from langfuse.api import DatasetStatus
-from tests.support.utils import create_uuid
+from tests.support.utils import create_uuid, wait_for_result
 
 
 def test_create_and_get_dataset():
@@ -98,7 +98,10 @@ def test_upsert_and_get_dataset_item():
     )
 
     # Instead, get all dataset items and find the one with matching ID
-    dataset = langfuse.get_dataset(name)
+    dataset = wait_for_result(
+        lambda: langfuse.get_dataset(name),
+        is_result_ready=lambda dataset: any(i.id == item.id for i in dataset.items),
+    )
     get_item = None
     for i in dataset.items:
         if i.id == item.id:
@@ -120,7 +123,16 @@ def test_upsert_and_get_dataset_item():
     )
 
     # Refresh dataset and find updated item
-    dataset = langfuse.get_dataset(name)
+    dataset = wait_for_result(
+        lambda: langfuse.get_dataset(name),
+        is_result_ready=lambda dataset: any(
+            i.id == item.id
+            and i.input == new_input
+            and i.expected_output == new_input
+            and i.status == DatasetStatus.ARCHIVED
+            for i in dataset.items
+        ),
+    )
     get_new_item = None
     for i in dataset.items:
         if i.id == item.id:

tests/live_provider/test_langchain.py

@@ -100,12 +100,13 @@ def test_callback_generated_from_lcel_chain():
 
     assert len(trace.observations) > 0
 
-    langchain_generation_span = list(
-        filter(
-            lambda o: o.type == "GENERATION" and o.name == "ChatOpenAI",
-            trace.observations,
-        )
-    )[0]
+    generation_observations = [
+        observation
+        for observation in trace.observations
+        if observation.type == "GENERATION"
+    ]
+    assert len(generation_observations) > 0
+    langchain_generation_span = generation_observations[0]
     langchain_root_spans = [
         observation
         for observation in trace.observations