Merge pull request #66 from meghanrjones/3dtest-update

Test against expected batch in test_batch_3d_2d_input

Author: Joe Hamman

xbatcher/tests/test_generators.py

@@ -133,7 +133,6 @@ def test_batch_1d_overlap(sample_ds_1d, olap):
 
 @pytest.mark.parametrize('bsize', [5, 10])
 def test_batch_3d_1d_input(sample_ds_3d, bsize):
-
     # first do the iteration over just one dimension
     bg = BatchGenerator(sample_ds_3d, input_dims={'x': bsize})
     for n, ds_batch in enumerate(bg):
@@ -164,8 +163,19 @@ def test_batch_3d_2d_input(sample_ds_3d, bsize):
         assert isinstance(ds_batch, xr.Dataset)
         assert ds_batch.dims['x'] == xbsize
         assert ds_batch.dims['y'] == bsize
-        # TODO? Is it worth it to try to reproduce the internal logic of the
-        # generator and verify that the slices are correct?
+        yn, xn = np.unravel_index(
+            n,
+            (
+                (sample_ds_3d.dims['y'] // bsize),
+                (sample_ds_3d.dims['x'] // xbsize),
+            ),
+        )
+        expected_xslice = slice(xbsize * xn, xbsize * (xn + 1))
+        expected_yslice = slice(bsize * yn, bsize * (yn + 1))
+        ds_batch_expected = sample_ds_3d.isel(
+            x=expected_xslice, y=expected_yslice
+        )
+        xr.testing.assert_equal(ds_batch_expected, ds_batch)
     assert (n + 1) == (
         (sample_ds_3d.dims['x'] // xbsize) * (sample_ds_3d.dims['y'] // bsize)
     )