Move streaming test to correct module

pola-rs · Jan 24, 2024 · e534370 · e534370
1 parent 25537c2
commit e534370
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 15 deletions.
diff --git a/py-polars/tests/unit/io/test_lazy_json.py b/py-polars/tests/unit/io/test_lazy_json.py
@@ -141,18 +141,3 @@ def test_anonymous_scan_explain(io_files_path: Path) -> None:
     assert "Anonymous" in q.explain()
     assert "Anonymous" in q.show_graph(raw_output=True)  # type: ignore[operator]
 
-
-def test_sink_ndjson_should_write_same_data(
-    io_files_path: Path, tmp_path: Path
-) -> None:
-    tmp_path.mkdir(exist_ok=True)
-    # Arrange
-    source_path = io_files_path / "foods1.csv"
-    target_path = tmp_path / "foods_test.ndjson"
-    expected = pl.read_csv(source_path)
-    lf = pl.scan_csv(source_path)
-    # Act
-    lf.sink_ndjson(target_path)
-    df = pl.read_ndjson(target_path)
-    # Assert
-    assert_frame_equal(df, expected)
diff --git a/py-polars/tests/unit/streaming/test_streaming_io.py b/py-polars/tests/unit/streaming/test_streaming_io.py
@@ -198,3 +198,21 @@ def test_streaming_cross_join_schema(tmp_path: Path) -> None:
     a.join(b, how="cross").sink_parquet(file_path)
     read = pl.read_parquet(file_path, parallel="none")
     assert read.to_dict(as_series=False) == {"a": [1, 2], "b": ["b", "b"]}
+
+
+@pytest.mark.write_disk()
+def test_sink_ndjson_should_write_same_data(
+    io_files_path: Path, tmp_path: Path
+) -> None:
+    tmp_path.mkdir(exist_ok=True)
+
+    source_path = io_files_path / "foods1.csv"
+    target_path = tmp_path / "foods_test.ndjson"
+
+    expected = pl.read_csv(source_path)
+
+    lf = pl.scan_csv(source_path)
+    lf.sink_ndjson(target_path)
+    df = pl.read_ndjson(target_path)
+
+    assert_frame_equal(df, expected)