Add more tests

setup.cfg

@@ -7,7 +7,7 @@ select = B,C,E,F,W,T4,B9
 
 [isort]
 known_first_party=xpublish
-known_third_party=dask,fastapi,numcodecs,numpy,pkg_resources,pytest,setuptools,starlette,uvicorn,xarray,zarr
+known_third_party=dask,fastapi,numcodecs,numpy,pandas,pkg_resources,pytest,setuptools,starlette,uvicorn,xarray,zarr
 multi_line_output=3
 include_trailing_comma=True
 force_grid_wrap=0

tests/test_zarr_compat.py

@@ -5,26 +5,161 @@
 
 import xpublish  # noqa: F401
 
-from .utils import TestMapper
+from .utils import TestMapper, create_dataset
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope='module')
 def airtemp_ds():
-    ds = xr.tutorial.open_dataset("air_temperature")
+    ds = xr.tutorial.open_dataset('air_temperature')
     return ds.chunk(dict(ds.dims))
 
 
-def test_zmetadata_identical(airtemp_ds):
+@pytest.mark.parametrize(
+    'start, end, freq, nlats, nlons, var_const, calendar, use_cftime',
+    [
+        ('2018-01-01', '2021-01-01', 'MS', 180, 360, True, 'standard', False),
+        ('2018-01-01', '2021-01-01', 'D', 180, 360, False, 'noleap', True),
+        ('2018-01-01', '2021-01-01', '6H', 180, 360, True, 'gregorian', False),
+        ('2018-01-01', '2050-01-01', 'A', 180, 360, None, '360_day', True),
+    ],
+)
+def test_zmetadata_identical(start, end, freq, nlats, nlons, var_const, calendar, use_cftime):
+    ds = create_dataset(
+        start=start,
+        end=end,
+        nlats=nlats,
+        nlons=nlons,
+        var_const=var_const,
+        use_cftime=use_cftime,
+        calendar=calendar,
+    )
+
+    ds = ds.chunk(ds.dims)
     zarr_dict = {}
-    airtemp_ds.to_zarr(zarr_dict, consolidated=True)
-    mapper = TestMapper(airtemp_ds.rest.app)
-    actual = json.loads(mapper[".zmetadata"].decode())
-    expected = json.loads(zarr_dict[".zmetadata"].decode())
+    ds.to_zarr(zarr_dict, consolidated=True)
+    mapper = TestMapper(ds.rest.app)
+    actual = json.loads(mapper['.zmetadata'].decode())
+    expected = json.loads(zarr_dict['.zmetadata'].decode())
     assert actual == expected
 
 
-def test_roundtrip(airtemp_ds):
-    mapper = TestMapper(airtemp_ds.rest.app)
+@pytest.mark.parametrize(
+    'start, end, freq, nlats, nlons, var_const, calendar, use_cftime',
+    [
+        ('2018-01-01', '2021-01-01', 'MS', 180, 360, True, 'standard', False),
+        ('2018-01-01', '2021-01-01', 'D', 180, 360, False, 'noleap', True),
+        ('2018-01-01', '2021-01-01', '6H', 180, 360, True, 'gregorian', False),
+        ('2018-01-01', '2050-01-01', 'A', 180, 360, None, '360_day', True),
+    ],
+)
+def test_roundtrip(start, end, freq, nlats, nlons, var_const, calendar, use_cftime):
+    ds = create_dataset(
+        start=start,
+        end=end,
+        nlats=nlats,
+        nlons=nlons,
+        var_const=var_const,
+        use_cftime=use_cftime,
+        calendar=calendar,
+    )
+    ds = ds.chunk(ds.dims)
+
+    mapper = TestMapper(ds.rest.app)
     actual = xr.open_zarr(mapper, consolidated=True)
 
-    xr.testing.assert_identical(actual, airtemp_ds)
+    xr.testing.assert_identical(actual, ds)
+
+
+xfail_reason = """Currently, xarray casts datetimes arrays to NumPy compatible arrays.
+This ends up producing unexpected behavior when calling encode_zarr_varible()
+on datasets with variables containing datetime like dtypes.
+
+See: https://github.com/jhamman/xpublish/pull/10#discussion_r388028417"""
+
+
+@pytest.mark.parametrize(
+    'start, end, freq, nlats, nlons, var_const, calendar, use_cftime, chunks, decode_times',
+    [
+        ('2018-01-01', '2021-01-01', 'MS', 180, 360, True, 'standard', False, {'time': 10}, False),
+        (
+            '2018-01-01',
+            '2021-01-01',
+            'D',
+            300,
+            600,
+            False,
+            'noleap',
+            True,
+            {'time': 10, 'lat': 300, 'lon': 300},
+            False,
+        ),
+        pytest.param(
+            '2018-01-01',
+            '2021-01-01',
+            'D',
+            300,
+            600,
+            False,
+            'noleap',
+            True,
+            {'time': 10, 'lat': 300, 'lon': 300},
+            True,
+            marks=pytest.mark.xfail(reason=xfail_reason),
+        ),
+        (
+            '2018-01-01',
+            '2021-01-01',
+            '12H',
+            180,
+            360,
+            True,
+            'gregorian',
+            False,
+            {'time': 36, 'lat': 10},
+            False,
+        ),
+        (
+            '2018-01-01',
+            '2038-01-01',
+            'A',
+            300,
+            600,
+            None,
+            '360_day',
+            True,
+            {'time': 10, 'lat': 75, 'lon': 120},
+            False,
+        ),
+        pytest.param(
+            '2018-01-01',
+            '2038-01-01',
+            'A',
+            300,
+            600,
+            None,
+            '360_day',
+            True,
+            {'time': 10, 'lat': 75, 'lon': 120},
+            True,
+            marks=pytest.mark.xfail(reason=xfail_reason),
+        ),
+    ],
+)
+def test_roundtrip_custom_chunks(
+    start, end, freq, nlats, nlons, var_const, calendar, use_cftime, chunks, decode_times
+):
+    ds = create_dataset(
+        start=start,
+        end=end,
+        nlats=nlats,
+        nlons=nlons,
+        var_const=var_const,
+        use_cftime=use_cftime,
+        calendar=calendar,
+        decode_times=decode_times,
+    )
+    ds = ds.chunk(chunks)
+    mapper = TestMapper(ds.rest.app)
+    actual = xr.open_zarr(mapper, consolidated=True, decode_times=decode_times)
+
+    xr.testing.assert_identical(actual, ds)

tests/utils.py

@@ -1,5 +1,13 @@
+from functools import reduce
+from operator import mul
+
+import numpy as np
+import pandas as pd
+import xarray as xr
 from starlette.testclient import TestClient
 
+rs = np.random.RandomState(np.random.MT19937(np.random.SeedSequence(123456789)))
+
 
 class TestMapper(TestClient):
     """
@@ -9,5 +17,99 @@ class TestMapper(TestClient):
     def __getitem__(self, key):
         response = self.get(key)
         if response.status_code != 200:
-            raise KeyError("{} not found. status_code = {}".format(key, response.status_code))
+            raise KeyError('{} not found. status_code = {}'.format(key, response.status_code))
         return response.content
+
+
+def create_dataset(
+    start='2018-01',
+    end='2020-12',
+    freq='MS',
+    calendar='standard',
+    units='days since 1980-01-01',
+    use_cftime=True,
+    decode_times=True,
+    nlats=1,
+    nlons=1,
+    var_const=None,
+):
+    """ Utility function for creating test data """
+
+    if use_cftime:
+        end = xr.coding.cftime_offsets.to_cftime_datetime(end, calendar=calendar)
+        dates = xr.cftime_range(start=start, end=end, freq=freq, calendar=calendar)
+
+    else:
+        dates = pd.date_range(start=pd.to_datetime(start), end=pd.to_datetime(end), freq=freq)
+
+    decoded_time_bounds = np.vstack((dates[:-1], dates[1:])).T
+
+    encoded_time_bounds = xr.coding.times.encode_cf_datetime(
+        decoded_time_bounds, units=units, calendar=calendar
+    )[0]
+
+    encoded_times = xr.DataArray(
+        encoded_time_bounds.mean(axis=1),
+        dims=('time'),
+        name='time',
+        attrs={'units': units, 'calendar': calendar},
+    )
+
+    decoded_times = xr.DataArray(
+        xr.coding.times.decode_cf_datetime(
+            encoded_times, units=units, calendar=calendar, use_cftime=use_cftime
+        ),
+        dims=['time'],
+    )
+    decoded_time_bounds = xr.DataArray(
+        decoded_time_bounds, name='time_bounds', dims=('time', 'd2'), coords={'time': decoded_times}
+    )
+
+    if decode_times:
+        times = decoded_times
+        time_bounds = decoded_time_bounds
+
+    else:
+        times = encoded_times
+        time_bounds = xr.DataArray(
+            encoded_time_bounds, name='time_bounds', dims=('time', 'd2'), coords={'time': times}
+        )
+
+    lats = np.linspace(start=-90, stop=90, num=nlats, dtype='float32')
+    lons = np.linspace(start=-180, stop=180, num=nlons, dtype='float32')
+
+    shape = (times.size, lats.size, lons.size)
+    num = reduce(mul, shape)
+
+    if var_const is None:
+        annual_cycle = np.sin(2 * np.pi * (decoded_times.dt.dayofyear.values / 365.25 - 0.28))
+        base = 10 + 15 * annual_cycle.reshape(-1, 1)
+        tmin_values = base + 3 * np.random.randn(annual_cycle.size, nlats * nlons)
+        tmax_values = base + 10 + 3 * np.random.randn(annual_cycle.size, nlats * nlons)
+        tmin_values = tmin_values.reshape(shape)
+        tmax_values = tmax_values.reshape(shape)
+
+    elif var_const:
+        tmin_values = np.ones(shape=shape)
+        tmax_values = np.ones(shape=shape) + 2
+
+    else:
+        tmin_values = np.arange(1, num + 1).reshape(shape)
+        tmax_values = np.arange(1, num + 1).reshape(shape)
+
+    ds = xr.Dataset(
+        {
+            'tmin': (('time', 'lat', 'lon'), tmin_values.astype('float32')),
+            'tmax': (('time', 'lat', 'lon'), tmax_values.astype('float32')),
+            'time_bounds': (('time', 'd2'), time_bounds),
+        },
+        {'time': times, 'lat': lats, 'lon': lons},
+    )
+
+    ds.tmin.encoding['_FillValue'] = np.float32(-9999999)
+    ds.tmax.encoding['_FillValue'] = np.float32(-9999999)
+    ds.time.encoding['bounds'] = 'time_bounds'
+    ds.time.encoding['units'] = units
+    ds.time.encoding['calendar'] = calendar
+
+    return ds