Coverage for src/rechunk_data/tests/conftest.py: 100%
59 statements
« prev ^ index » next coverage.py v7.3.1, created at 2023-09-30 09:58 +0000
« prev ^ index » next coverage.py v7.3.1, created at 2023-09-30 09:58 +0000
1"""pytest definitions to run the unittests."""
3from pathlib import Path
4from tempfile import TemporaryDirectory, NamedTemporaryFile
5from typing import Generator, Tuple
7import dask
8import pytest
9import numpy as np
10import xarray as xr
13def create_data(
14 variable_name: str, chunk_size: Tuple[int, ...], dims: Tuple[str, ...]
15) -> xr.Dataset:
16 """Create a netcdf dataset."""
17 coords = {d: np.ones(chunk_size[n]) for (n, d) in enumerate(dims)}
18 dset = xr.DataArray(
19 np.zeros(chunk_size),
20 dims=dims,
21 coords=coords,
22 name=variable_name,
23 ).chunk(dict(zip(dims, chunk_size)))
24 dset.encoding = {"chunksizes": chunk_size}
25 return xr.Dataset({variable_name: dset})
28@pytest.fixture(scope="session")
29def small_chunk() -> Generator[Tuple[int, int, int, int], None, None]:
30 """Define tuple for smaller chunks sizes."""
31 yield (1, 1, 24, 24)
34@pytest.fixture(scope="session")
35def empty_data() -> Generator[xr.Dataset, None, None]:
36 """Create an empty datasset."""
38 yield xr.Dataset(
39 {"tas_bnds": xr.DataArray(["hallo"], name="tas_bnds", dims=("lon",))}
40 )
43@pytest.fixture(scope="session")
44def large_chunk() -> Generator[Tuple[int, int, int, int], None, None]:
45 """Define tuple for smaller chunks sizes."""
46 yield (720, 12, 4, 4)
49@pytest.fixture(scope="session")
50def dims() -> Generator[Tuple[str, str, str, str], None, None]:
51 """Dimensions of all datasets."""
52 yield ("time", "height", "Latitude", "Longitude")
55@pytest.fixture(scope="function")
56def temp_dir() -> Generator[Path, None, None]:
57 """Temporary Directory for creating data files."""
58 with TemporaryDirectory() as temporary_dir:
59 yield Path(temporary_dir)
62@pytest.fixture(scope="session")
63def small_chunk_data(
64 small_chunk: Tuple[int, int, int, int],
65 dims: Tuple[str, str, str, str],
66 variable_name: str,
67) -> Generator[xr.Dataset, None, None]:
68 """Create a dataset with small chunks."""
69 yield create_data(variable_name, small_chunk, dims)
72@pytest.fixture(scope="function")
73def large_chunk_data(
74 large_chunk: Tuple[int, int, int, int],
75 dims: Tuple[str, str, str, str],
76 variable_name: str,
77) -> Generator[xr.Dataset, None, None]:
78 """Create a dataset with small chunks."""
79 with dask.config.set({"array.chunk-size": "1MiB"}):
80 yield create_data(variable_name, large_chunk, dims)
83@pytest.fixture(scope="session")
84def variable_name() -> str:
85 return "tas"
88@pytest.fixture(scope="function")
89def data_dir(
90 temp_dir: Path,
91 variable_name: str,
92 small_chunk: Tuple[int, int, int, int],
93 dims: Tuple[str, str, str],
94 small_chunk_data: xr.Dataset,
95) -> Generator[Path, None, None]:
96 """Create a directory with netcdf files."""
97 encoding = {variable_name: {"chunksizes": small_chunk}}
98 for number in range(1, 10):
99 file_name = temp_dir / "foo" / "bar" / f"tas_model1_{number}.nc"
100 file_name.parent.mkdir(parents=True, exist_ok=True)
101 small_chunk_data.to_netcdf(file_name, encoding=encoding)
102 yield temp_dir
105@pytest.fixture(scope="function")
106def data_file(
107 variable_name: str,
108 large_chunk: Tuple[int, int, int, int],
109 dims: Tuple[str, str, str],
110 large_chunk_data: xr.Dataset,
111) -> Generator[Path, None, None]:
112 """Create a directory with netcdf files."""
113 encoding = {variable_name: {"chunksizes": large_chunk}}
114 with NamedTemporaryFile(suffix=".nc") as temp_file:
115 file_name = Path(temp_file.name)
116 large_chunk_data.to_netcdf(file_name, encoding=encoding)
117 yield file_name
120@pytest.fixture(scope="function")
121def wrong_file_type(temp_dir) -> Generator[Path, None, None]:
122 """Temporary Directory contaning non netcdf files."""
124 for _ in range(1, 3):
125 file_name = temp_dir / "worng_file.txt"
126 file_name.touch()
127 yield temp_dir