fix failing test

Author: Joseph Hamman

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v3.1.0
+    rev: v4.0.1
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
@@ -10,25 +10,26 @@ repos:
       - id: double-quote-string-fixer
 
   - repo: https://github.com/ambv/black
-    rev: 19.10b0
+    rev: 21.5b2
     hooks:
       - id: black
         args: ["--line-length", "80", "--skip-string-normalization"]
 
   - repo: https://gitlab.com/pycqa/flake8
-    rev: 3.8.3
+    rev: 3.9.2
     hooks:
       - id: flake8
   - repo: https://github.com/asottile/seed-isort-config
     rev: v2.2.0
     hooks:
       - id: seed-isort-config
   - repo: https://github.com/pre-commit/mirrors-isort
-    rev: v5.2.1
+    rev: v5.8.0
     hooks:
       - id: isort
 
-  - repo: https://github.com/prettier/prettier
-    rev: 2.0.5
+  - repo: https://github.com/pre-commit/mirrors-prettier
+    rev: v2.2.0
     hooks:
       - id: prettier
+        language_version: system

conftest.py

@@ -0,0 +1,9 @@
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def add_standard_imports(doctest_namespace, tmpdir):
+    import numpy as np
+
+    # always seed numpy.random to make the examples deterministic
+    np.random.seed(0)

xbatcher/generators.py

@@ -49,7 +49,7 @@ def _drop_input_dims(ds, input_dims, suffix='_input'):
         out = out.rename({dim: newdim})
         # extra steps needed if there is a coordinate
         if newdim in out:
-            out = out.drop(newdim)
+            out = out.drop_vars(newdim)
             out.coords[dim] = newdim, ds[dim].data, ds[dim].attrs
     return out
 

xbatcher/tests/test_generators.py

@@ -47,7 +47,7 @@ def test_batch_1d_concat(sample_ds_1d, bsize):
 @pytest.mark.parametrize('bsize', [5, 10])
 def test_batch_1d_no_coordinate(sample_ds_1d, bsize):
     # fix for #3
-    ds_dropped = sample_ds_1d.drop('x')
+    ds_dropped = sample_ds_1d.drop_vars('x')
     bg = BatchGenerator(ds_dropped, input_dims={'x': bsize})
     for n, ds_batch in enumerate(bg):
         assert isinstance(ds_batch, xr.Dataset)
@@ -60,7 +60,7 @@ def test_batch_1d_no_coordinate(sample_ds_1d, bsize):
 @pytest.mark.parametrize('bsize', [5, 10])
 def test_batch_1d_concat_no_coordinate(sample_ds_1d, bsize):
     # test for #3
-    ds_dropped = sample_ds_1d.drop('x')
+    ds_dropped = sample_ds_1d.drop_vars('x')
     bg = BatchGenerator(
         ds_dropped, input_dims={'x': bsize}, concat_input_dims=True
     )
@@ -118,7 +118,7 @@ def test_batch_3d_1d_input(sample_ds_3d, bsize):
         expected_slice = slice(bsize * n, bsize * (n + 1))
         ds_batch_expected = (
             sample_ds_3d.isel(x=expected_slice)
-            .stack(sample=['y', 'time'])
+            .stack(sample=['time', 'y'])
             .transpose('sample', 'x')
         )
         print(ds_batch)