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

1"""pytest definitions to run the unittests.""" 

2 

3from pathlib import Path 

4from tempfile import TemporaryDirectory, NamedTemporaryFile 

5from typing import Generator, Tuple 

6 

7import dask 

8import pytest 

9import numpy as np 

10import xarray as xr 

11 

12 

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}) 

26 

27 

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) 

32 

33 

34@pytest.fixture(scope="session") 

35def empty_data() -> Generator[xr.Dataset, None, None]: 

36 """Create an empty datasset.""" 

37 

38 yield xr.Dataset( 

39 {"tas_bnds": xr.DataArray(["hallo"], name="tas_bnds", dims=("lon",))} 

40 ) 

41 

42 

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) 

47 

48 

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") 

53 

54 

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) 

60 

61 

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) 

70 

71 

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) 

81 

82 

83@pytest.fixture(scope="session") 

84def variable_name() -> str: 

85 return "tas" 

86 

87 

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 

103 

104 

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 

118 

119 

120@pytest.fixture(scope="function") 

121def wrong_file_type(temp_dir) -> Generator[Path, None, None]: 

122 """Temporary Directory contaning non netcdf files.""" 

123 

124 for _ in range(1, 3): 

125 file_name = temp_dir / "worng_file.txt" 

126 file_name.touch() 

127 yield temp_dir