ENH: add tests

Author: kmuehlbauer

xarray/tests/test_concat.py

@@ -571,6 +571,57 @@ def test_concat_str_dtype(self, dtype, dim) -> None:
 
         assert np.issubdtype(actual.x2.dtype, dtype)
 
+    @pytest.mark.parametrize("dim", [True, False])
+    @pytest.mark.parametrize("coord", [True, False])
+    def test_concat_fill_missing_variables(self, dim, coord):
+        # create var names list with one missing value
+        def get_var_names(var_cnt=10, list_cnt=10):
+            orig = [f'd{i:02d}' for i in range(var_cnt)]
+            var_names = []
+            for i in range(0, list_cnt):
+                l1 = orig.copy()
+                var_names.append(l1)
+            return var_names
+
+        def create_ds(var_names, dim=False, coord=False, drop_idx=False):
+            out_ds = []
+            ds = Dataset()
+            ds = ds.assign_coords({"x": np.arange(2)})
+            ds = ds.assign_coords({"y": np.arange(3)})
+            ds = ds.assign_coords({"z": np.arange(4)})
+            for i, dsl in enumerate(var_names):
+                vlist = dsl.copy()
+                if drop_idx:
+                    vlist.pop(drop_idx[i])
+                foo_data = np.arange(48, dtype=float).reshape(2, 2, 3, 4)
+                dsi = ds.copy()
+                if coord:
+                    dsi = ds.assign({"time": (["time"], [i * 2, i * 2 + 1])})
+                for k in vlist:
+                    dsi = dsi.assign({k: (["time", "x", "y", "z"], foo_data.copy())})
+                if not dim:
+                    dsi = dsi.isel(time=0)
+                out_ds.append(dsi)
+            return out_ds
+        var_names = get_var_names()
+
+        import random
+        random.seed(42)
+        drop_idx = [random.randrange(len(vlist)) for vlist in var_names]
+        expected = concat(create_ds(var_names, dim=dim, coord=coord), dim="time", data_vars="all")
+        for i, idx in enumerate(drop_idx):
+            if dim:
+                expected[var_names[0][idx]][i * 2: i * 2 + 2] = np.nan
+            else:
+                expected[var_names[0][idx]][i] = np.nan
+
+        concat_ds = create_ds(var_names, dim=dim, coord=coord, drop_idx=drop_idx)
+        actual = concat(concat_ds, dim="time", data_vars="all")
+
+        for name in var_names[0]:
+            assert_equal(expected[name], actual[name])
+        assert_equal(expected, actual)
+
 
 class TestConcatDataArray:
     def test_concat(self) -> None: