push

Author: hassiebp

tests/test_datasets.py

@@ -7,8 +7,7 @@
 from langchain_openai import OpenAI
 
 from langfuse import Langfuse, observe
-from langfuse.api import DatasetStatus
-from langfuse.api import Observation
+from langfuse.api import DatasetStatus, Observation
 from langfuse.langchain import CallbackHandler
 from tests.utils import create_uuid, get_api
 
@@ -220,7 +219,7 @@ def test_get_dataset_runs():
 
     langfuse.flush()
     time.sleep(1)  # Give API time to process
-    runs = langfuse.api
+    runs = langfuse.api.datasets.get_runs(dataset_name=dataset_name)
 
     assert len(runs.data) == 2
     assert runs.data[0].name == run_name_2

tests/test_propagate_attributes.py

@@ -7,13 +7,16 @@
 
 import concurrent.futures
 import time
+from datetime import datetime
 
 import pytest
 from opentelemetry.instrumentation.threading import ThreadingInstrumentor
 
 from langfuse import propagate_attributes
 from langfuse._client.attributes import LangfuseOtelSpanAttributes, _serialize
 from langfuse._client.constants import LANGFUSE_SDK_EXPERIMENT_ENVIRONMENT
+from langfuse._client.datasets import DatasetClient, DatasetItemClient
+from langfuse.api import Dataset, DatasetItem, DatasetStatus
 from tests.test_otel import TestOTelBase
 
 
@@ -2437,12 +2440,6 @@ def test_experiment_attributes_propagate_with_dataset(
         self, langfuse_client, memory_exporter, monkeypatch
     ):
         """Test experiment attribute propagation with Langfuse dataset."""
-        import time
-        from datetime import datetime
-
-        from langfuse._client.attributes import _serialize
-        from langfuse._client.datasets import DatasetClient, DatasetItemClient
-        from langfuse.model import Dataset, DatasetItem, DatasetStatus
 
         # Mock the async API to create dataset run items
         async def mock_create_dataset_run_item(*args, **kwargs):