Test halo accessor

Author: cphyc

yt/frontends/rockstar/data_structures.py

@@ -1,7 +1,7 @@
 import glob
 import os
 from functools import cached_property
-from typing import Optional
+from typing import List, Optional
 
 import numpy as np
 
@@ -21,7 +21,7 @@ class RockstarBinaryFile(HaloCatalogFile):
     _position_offset: int
     _member_offset: int
     _Npart: np.array
-    _ids_halos: list[int]
+    _ids_halos: List[int]
     _file_size: int
 
     def __init__(self, ds, io, filename, file_id, range):

yt/frontends/rockstar/tests/test_outputs.py

@@ -38,3 +38,23 @@ def test_particle_selection():
     ds = data_dir_load(r1)
     psc = ParticleSelectionComparison(ds)
     psc.run_defaults()
+
+
+@requires_file(r1)
+def test_halo_loading():
+    ds = data_dir_load(r1)
+
+    for halo_id, Npart in zip(
+        ds.r["halos", "particle_identifier"],
+        ds.r["halos", "num_p"],
+    ):
+        halo = ds.halo("halos", halo_id)
+        assert halo is not None
+
+        # Try accessing properties
+        halo.position
+        halo.velocity
+        halo.mass
+
+        # Make sure we can access the member particles
+        assert len(halo.member_ids) == Npart