better test?

DeaMariaLeon · Oct 28, 2024 · e9e8c28 · e9e8c28
1 parent 892aa2c
commit e9e8c28
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/tests/expr_and_series/map_batches_test.py b/tests/expr_and_series/map_batches_test.py
@@ -9,22 +9,22 @@
 
 data = {"a": [1, 2, 3], "b": [4, 5, 6], "z": [7.0, 8.0, 9.0]}
 
-expected = [2, 3, 4]
-
 
 def test_map_batches_expr(constructor: Constructor) -> None:
     if is_dask_dataframe(constructor(data)):  # Remove
         pytest.skip()
     df = nw.from_native(constructor(data))
-    e = df.select(nw.col("a").map_batches(lambda s: s + 1))
-    assert_equal_data(e, {"a": expected})
+    e = df.select(nw.col("a", "b").map_batches(lambda s: s + 1))
+    assert_equal_data(e, {"a": [2, 3, 4], "b": [5, 6, 7]})
 
 
 def test_map_batches_expr_numpy(constructor: Constructor) -> None:
     if is_dask_dataframe(constructor(data)):  # Remove
         pytest.skip()
     df = nw.from_native(constructor(data))
     e = df.select(
-        nw.col("a").map_batches(lambda s: s.to_numpy() + 1, return_dtype=nw.Float64).sum()
+        nw.col("a")
+        .map_batches(lambda s: s.to_numpy() + 1, return_dtype=nw.Float64())
+        .sum()
     )
     assert_equal_data(e, {"a": [9.0]})