test: add tests for qsel.average

tests/plotting/test_general.py

@@ -213,3 +213,59 @@ def test_plot_array(data, colorbar, xlim, ylim, crop, kwargs) -> None:
         else:
             assert ax.get_ylim() == ylim
     plt.close()
+
+
+def test_qsel_average_single_dim() -> None:
+    # Create a simple 2D DataArray with dims 'x' and 'y'
+
+    x = np.array([10, 20])
+    y = np.array([30, 40, 50])
+    data = np.array([[1, 2, 3], [4, 5, 6]])
+    da = xr.DataArray(data, dims=("x", "y"), coords={"x": x, "y": y})
+
+    # Average over the 'x' dimension using qsel.average.
+    # The expected result is the mean along 'x' and retaining the averaged coordinate.
+    # Expected mean data: [[(1+4)/2, (2+5)/2, (3+6)/2]] = [[2.5, 3.5, 4.5]]
+    expected = data.mean(axis=0)
+    result = da.qsel.average("x")
+
+    # After averaging, 'x' should not be a dimension; it is stored as a coordinate.
+    assert "x" not in result.dims
+    # Compare the resulting data with the expected average.
+    np.testing.assert_allclose(result.data, expected)
+    # Check that the coordinate 'x' is the mean of the original x-values.
+    np.testing.assert_allclose(result.coords["x"].data, x.mean())
+
+
+def test_qsel_average_multiple_dim() -> None:
+    # Create a simple 2D DataArray with dims 'x' and 'y'
+
+    x = np.array([0, 10, 20])
+    y = np.array([100, 200])
+    data = np.array([[1, 2], [4, 5], [7, 8]])
+    da = xr.DataArray(data, dims=("x", "y"), coords={"x": x, "y": y})
+
+    # Average over both 'x' and 'y'
+    # Expected result is a scalar value: mean of all data.
+    expected = data.mean()
+    result = da.qsel.average(["x", "y"])
+
+    # The resulting DataArray should have no dimensions.
+    assert not result.dims
+    # Data should be equal to the overall mean.
+    np.testing.assert_allclose(result.data, expected)
+    # Coordinates are retained as scalars equal to the mean of the original coords.
+    np.testing.assert_allclose(result.coords["x"].data, x.mean())
+    np.testing.assert_allclose(result.coords["y"].data, y.mean())
+
+
+def test_qsel_average_invalid_dim() -> None:
+    # Create a simple 1D DataArray
+
+    x = np.array([0, 1, 2])
+    data = np.array([10, 20, 30])
+    da = xr.DataArray(data, dims=("x",), coords={"x": x})
+
+    # Averaging over an invalid dimension should raise a ValueError
+    with pytest.raises(ValueError, match="Dimension `z` not found in data"):
+        _ = da.qsel.average("z")

tests/utils/test_array.py

@@ -32,7 +32,7 @@ def testfunc(x, y):
     expected = xr.DataArray(expected_vals, coords={"x": x_val, "y": y_val})
 
     xr.testing.assert_identical(
-        testfunc(
+        testfunc_(
             xr.DataArray(x_val, coords={"x": x_val}),
             xr.DataArray(y_val, coords={"y": y_val}),
         ),