Coverage for src/rechunk_data/tests/test_rechunking.py: 100%

1"""Unit tests for rechunking the data.""" 

2import dask 

3import xarray as xr 

4 

5from rechunk_data._rechunk import _rechunk_dataset 

6 

7 

8def test_rechunking_small_data( 

9 small_chunk_data: xr.Dataset, variable_name: str 

10) -> None: 

11 """Rechunking of small datasets should have no effect.""" 

12 chunks = small_chunk_data[variable_name].data.chunksize 

13 dset, encoding = _rechunk_dataset(small_chunk_data, "h5netcdf") 

14 assert dset[variable_name].data.chunksize == chunks 

15 dset, encoding = _rechunk_dataset(small_chunk_data.load(), "h5netcdf") 

16 assert encoding == {} 

17 

18 

19def test_empty_dataset(empty_data): 

20 """Test handling of empyt data.""" 

21 with dask.config.set({"array.chunk-size": "0.01b"}): 

22 dset, encoding = _rechunk_dataset(empty_data, "h5netcdf") 

23 assert encoding == {} 

24 

25 

26def test_rechunking_large_data( 

27 large_chunk_data: xr.Dataset, variable_name: str 

28) -> None: 

29 """Check the rechunk_data of a large dataset.""" 

30 chunks = ( 

31 large_chunk_data[variable_name] 

32 .data.rechunk({0: "auto", 1: "auto", 2: None, 3: None}) 

33 .chunksize 

34 ) 

35 dset, encoding = _rechunk_dataset(large_chunk_data, "h5netcdf") 

36 assert encoding[variable_name]["chunksizes"] == chunks 

37 assert dset[variable_name].data.chunksize == chunks