tests(ci): fix flaky tests

Author: hassiebp

tests/test_core_sdk.py

@@ -1479,6 +1479,7 @@ def mask_func(data):
     span.update(output="new_classified")
 
     langfuse.flush()
+    sleep(1)
 
     fetched_trace = api_wrapper.get_trace(trace.id)
     assert fetched_trace["input"] == {"sensitive": "MASKED"}
@@ -1503,6 +1504,7 @@ def faulty_mask_func(data):
     sleep(0.1)
     trace.update(output={"more": "sensitive"})
     langfuse.flush()
+    sleep(1)
 
     fetched_trace = api_wrapper.get_trace(trace.id)
     assert fetched_trace["input"] == "<fully masked due to failed mask function>"

tests/test_datasets.py

@@ -431,8 +431,9 @@ def test_llama_index_dataset():
             index.as_query_engine().query(
                 "What did the speaker achieve in the past twelve months?"
             )
-
     langfuse.flush()
+    handler.flush()
+
     run = langfuse.get_dataset_run(dataset_name, run_name)
 
     assert run.name == run_name