From 038436365d4757a322cac37307503c132d1fe2a7 Mon Sep 17 00:00:00 2001
From: Deepak Cherian <dcherian@users.noreply.github.com>
Date: Mon, 4 Nov 2024 20:54:19 -0700
Subject: [PATCH] Fix groupby tests (#9716)

Closes #9715
---
 xarray/tests/test_groupby.py | 39 ++++++++++++++++++++++--------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py
index 95ba0a3384e..c0eeace71af 100644
--- a/xarray/tests/test_groupby.py
+++ b/xarray/tests/test_groupby.py
@@ -2848,20 +2848,27 @@ def test_multiple_groupers(use_flox) -> None:
                 xy=UniqueGrouper(labels=["a", "b", "c"]),
                 eagerly_compute_group=eagerly_compute_group,
             )
-            with raise_if_dask_computes(max_computes=1):
-                if eagerly_compute_group:
-                    with pytest.warns(DeprecationWarning):
-                        gb = b.groupby(**kwargs)  # type: ignore[arg-type]
-                else:
-                    gb = b.groupby(**kwargs)  # type: ignore[arg-type]
-                    assert is_chunked_array(gb.encoded.codes.data)
-                    assert not gb.encoded.group_indices
             expected = xr.DataArray(
                 [[[1, 1, 1], [np.nan, 1, 2]]] * 4,
                 dims=("z", "x", "xy"),
                 coords={"xy": ("xy", ["a", "b", "c"], {"foo": "bar"})},
             )
-            assert_identical(gb.count(), expected)
+            if eagerly_compute_group:
+                with raise_if_dask_computes(max_computes=1):
+                    with pytest.warns(DeprecationWarning):
+                        gb = b.groupby(**kwargs)  # type: ignore[arg-type]
+                    assert_identical(gb.count(), expected)
+            else:
+                with raise_if_dask_computes(max_computes=0):
+                    gb = b.groupby(**kwargs)  # type: ignore[arg-type]
+                assert is_chunked_array(gb.encoded.codes.data)
+                assert not gb.encoded.group_indices
+                if has_flox:
+                    with raise_if_dask_computes(max_computes=1):
+                        assert_identical(gb.count(), expected)
+                else:
+                    with pytest.raises(ValueError, match="when lazily grouping"):
+                        gb.count()
 
 
 @pytest.mark.parametrize("use_flox", [True, False])
@@ -3003,11 +3010,6 @@ def test_lazy_grouping(grouper, expect_index):
     pd.testing.assert_index_equal(encoded.full_index, expect_index)
     np.testing.assert_array_equal(encoded.unique_coord.values, np.array(expect_index))
 
-    lazy = (
-        xr.Dataset({"foo": data}, coords={"zoo": data})
-        .groupby(zoo=grouper, eagerly_compute_group=False)
-        .count()
-    )
     eager = (
         xr.Dataset({"foo": data}, coords={"zoo": data.compute()})
         .groupby(zoo=grouper)
@@ -3017,9 +3019,16 @@ def test_lazy_grouping(grouper, expect_index):
         {"foo": (encoded.codes.name, np.ones(encoded.full_index.size))},
         coords={encoded.codes.name: expect_index},
     )
-    assert_identical(eager, lazy)
     assert_identical(eager, expected)
 
+    if has_flox:
+        lazy = (
+            xr.Dataset({"foo": data}, coords={"zoo": data})
+            .groupby(zoo=grouper, eagerly_compute_group=False)
+            .count()
+        )
+        assert_identical(eager, lazy)
+
 
 @requires_dask
 def test_lazy_grouping_errors():