Add a partitioning test

Author: AstroPatty

test/parallel/test_mpi.py

@@ -71,6 +71,21 @@ def test_mpi(input_path):
     parallel_assert(len(data) != 0)
 
 
+@pytest.mark.parallel(nprocs=4)
+def test_partitioning_includes_all(input_path):
+    with oc.open(input_path) as f:
+        tags = f.select("fof_halo_tag").data
+
+    comm = mpi4py.MPI.COMM_WORLD
+    all_tags = comm.allgather(tags)
+    all_tags = reduce(lambda left, right: left.union(set(right)), all_tags, set())
+
+    file = h5py.File(input_path)
+    original_tags = set(file["data"]["fof_halo_tag"][:])
+
+    parallel_assert(all_tags == original_tags)
+
+
 @pytest.mark.parallel(nprocs=4)
 def test_take(input_path):
     ds = oc.open(input_path)

test/test_diffsky.py

@@ -35,15 +35,6 @@ def test_comoving_to_unitless(core_path_487):
         assert np.all(data[col].value == data_unitless[col].value)
 
 
-def test_open_multiple(core_path_487, core_path_475):
-    ds = oc.open(core_path_487, core_path_475, synth_cores=True)
-    assert len(ds.keys()) == 4
-    ds = oc.open(core_path_487, core_path_475)
-    assert len(ds.keys()) == 2
-    z_range = ds.z_range
-    assert z_range[1] - z_range[0] > 0.05
-
-
 def test_open_multiple_write(core_path_487, core_path_475, tmp_path):
     ds = oc.open(core_path_487, core_path_475, synth_cores=True)
     original_length = len(ds)