add unit testing for sink_json method

pola-rs · Sep 6, 2023 · c574b9e · c574b9e
1 parent 49c1d37
commit c574b9e
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 12 deletions.
diff --git a/py-polars/src/conversion.rs b/py-polars/src/conversion.rs
@@ -1287,7 +1287,7 @@ impl FromPyObject<'_> for Wrap<JsonFormat> {
             "json_lines" => JsonFormat::JsonLines,
             v => {
                 return Err(PyValueError::new_err(format!(
-                    "json fommat must be one of: {{'json', 'json_lines'}}, got {v}",
+                    "json format must be one of: {{'json', 'json_lines'}}, got {v}",
                 )))
             },
         };

diff --git a/py-polars/tests/unit/io/test_lazy_json.py b/py-polars/tests/unit/io/test_lazy_json.py
@@ -108,24 +108,52 @@ def test_sink_json_should_write_same_data(io_files_path: Path, tmp_path: Path) -
     source_path = io_files_path / "foods1.csv"
     target_path = tmp_path / "foods_test.ndjson"
     expected = pl.read_csv(source_path)
-    lz_df = pl.scan_csv(source_path)
+    lf = pl.scan_csv(source_path)
     # Act
-    lz_df.sink_json(target_path)
+    lf.sink_json(target_path)
     df = pl.read_ndjson(target_path)
     # Assert
     assert_frame_equal(df, expected)
 
 
-def test_sink_json_should_support_with_options(io_files_path: Path, tmp_path: Path) -> None:
+def test_sink_json_should_write_same_data_with_json_argument(io_files_path: Path, tmp_path: Path) -> None:
     tmp_path.mkdir(exist_ok=True)
-
     # Arrange
-    source_path = io_files_path / "foods1.ndjson"
+    source_path = io_files_path / "foods1.csv"
     target_path = tmp_path / "foods_test.json"
-    expected = pl.read_ndjson(source_path)
-    lz_df = pl.scan_ndjson(source_path)
+
+    expected = pl.read_csv(source_path)
+    lf = pl.scan_csv(source_path)
     # Act
-    lz_df.sink_json(target_path)
-    # df = pl.read_json(target_path)
-    # # Assert
-    # assert_frame_equal(df, expected)
+    lf.sink_json(target_path,  json_format="json")
+    df = pl.read_json(target_path)
+    # Assert
+    assert_frame_equal(df, expected)
+
+
+def test_sink_json_should_write_same_data_with_json_lines_argument(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_json(target_path,  json_format="json_lines")
+    df = pl.read_ndjson(target_path)
+# Assert
+    assert_frame_equal(df, expected)
+
+
+def test_sink_json_should_raise_exception_with_invalid_argument(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"
+
+    lf = pl.scan_csv(source_path)
+    # Act & Assert
+    with pytest.raises(ValueError):
+        lf.sink_json(target_path,  json_format="invalid_argument")
+