Create python module

src/rechunk_data/tests/test_rechunk_netcdf.py

@@ -4,7 +4,8 @@ from tempfile import NamedTemporaryFile, TemporaryDirectory
 from pathlib import Path
 
 import dask
-from rechunk_data import rechunk_netcdf_file, _save_dataset
+from rechunk_data import rechunk_netcdf_file, rechunk_dataset
+from rechunk_data._rechunk import _save_dataset
 
 
 def test_rechunk_data_dir_with_overwrite(data_dir: Path) -> None:
@@ -38,6 +39,13 @@ def test_rechunk_single_data_file(data_file: Path) -> None:
             assert Path(temp_file.name).exists()
 
 
+def test_rechunk_dataset(small_chunk_data) -> None:
+    """Test rechunking an xarray dataset."""
+    with dask.config.set({"array.chunk-size": "1MiB"}):
+        new_data = rechunk_dataset(small_chunk_data)
+        assert list(new_data.data_vars) == list(small_chunk_data.data_vars)
+
+
 def test_wrong_or_format(small_chunk_data, caplog) -> None:
     """Testing wrong file format."""
     caplog.clear()