test

tests/integrations/pandas-test/test_pandas.py

@@ -27,10 +27,12 @@ def test_dataset(client):
     rows = [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}]
     ds = Dataset(rows=rows)
     df = ds.to_pandas()
+    assert df["a"].tolist() == [1, 3, 5]
+    assert df["b"].tolist() == [2, 4, 6]
 
     df2 = pd.DataFrame(rows)
     ds2 = Dataset.from_pandas(df2)
-
+    assert ds2.rows == rows
     assert df.equals(df2)
     assert ds.rows == ds2.rows
 
@@ -46,9 +48,10 @@ def greet(name: str, age: int) -> str:
     calls = greet.calls()
     dataset = Dataset.from_calls(calls)
     df = dataset.to_pandas()
-
-    assert df["name"].tolist() == ["Alice", "Bob"]
-    assert df["age"].tolist() == [30, 25]
+    assert df["inputs"].tolist() == [
+        {"name": "Alice", "age": 30},
+        {"name": "Bob", "age": 25},
+    ]
     assert df["output"].tolist() == [
         "Hello, Alice! You are 30 years old.",
         "Hello, Bob! You are 25 years old.",