add multi partition test

lincc-frameworks · Nov 14, 2023 · 55ee6b7 · 55ee6b7
1 parent ecc6e03
commit 55ee6b7
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tests/tape_tests/test_ensemble.py b/tests/tape_tests/test_ensemble.py
@@ -823,10 +823,14 @@ def test_keep_zeros(parquet_ensemble):
     ],
 )
 @pytest.mark.parametrize("by_band", [True, False])
-def test_calc_nobs(data_fixture, request, by_band):
+@pytest.mark.parametrize("multi_partition", [True, False])
+def test_calc_nobs(data_fixture, request, by_band, multi_partition):
     # Get the Ensemble from a fixture
     ens = request.getfixturevalue(data_fixture)
 
+    if multi_partition:
+        ens._source = ens._source.repartition(3)
+
     # Drop the existing nobs columns
     ens._object = ens._object.drop(["nobs_g", "nobs_r", "nobs_total"], axis=1)