Fix some tests broken by the changes

Author: AstroPatty

src/opencosmo/dataset/dataset.py

@@ -788,7 +788,9 @@ def collect(self) -> Dataset:
 
         If working in an MPI context, all ranks will recieve the same data.
         """
-        new_handler = self.__handler.collect(self.__state.columns, self.__state.index)
+        new_handler = self.__handler.collect(
+            self.__state.unit_handlers.keys(), self.__state.index
+        )
         new_index = ChunkedIndex.from_size(len(new_handler))
         new_state = self.__state.with_index(new_index)
         return Dataset(

src/opencosmo/dataset/state.py

@@ -388,6 +388,15 @@ def with_units(
         """
 
         convention_ = UnitConvention(convention)
+        if (
+            convention_ == UnitConvention.SCALEFREE
+            and UnitConvention(self.header.file.unit_convention)
+            != UnitConvention.SCALEFREE
+        ):
+            raise ValueError(
+                f"Cannot convert units with convention {self.header.file.unit_convention} to convention scalefree"
+            )
+
         return DatasetState(
             self.__unit_applicators,
             self.__index,

src/opencosmo/header.py

@@ -257,6 +257,7 @@ def read_header(
             f"Error: {e}"
         )
 
+    print(file_parameters.unit_convention)
     origin_parameter_models = origin.get_origin_parameters(file_parameters.origin)
     required_origin_params, optional_origin_params = read_origin_parameters(
         file, origin_parameter_models

src/opencosmo/parameters/file.py

@@ -59,9 +59,17 @@ def parse_region(cls, data):
     def validate_is_lightcone(cls, value):
         return bool(value)
 
+    @field_validator("unit_convention", mode="before")
+    def validate_convention(cls, value):
+        if isinstance(value, str):
+            return UnitConvention(value)
+        return value
+
     @field_serializer("unit_convention")
     def serialize_convention(self, value):
-        return value.value
+        if isinstance(value, UnitConvention):
+            return value.value
+        return value
 
     @model_serializer(mode="wrap")
     def serialize_model(self, handle):

src/opencosmo/units/converters.py

@@ -183,7 +183,7 @@ def physical_to_scalefree(
     return add_littleh(new_value, cosmology, base_unit)
 
 
-def raise_convert_error(from_: UnitConvention, to_: UnitConvention, *args, **kwargs):
+def raise_convert_error(*args, from_: UnitConvention, to_: UnitConvention, **kwargs):
     raise ValueError(
         f"Units in convention {str(from_)} cannot be converted to units in convention {str(to_)}"
     )

test/test_diffsky.py

@@ -23,7 +23,7 @@ def test_comoving_to_physical(core_path_487):
 
 
 def test_comoving_to_scalefree(core_path_487):
-    with pytest.raises(oc.transformations.units.UnitError):
+    with pytest.raises(ValueError):
         _ = oc.open(core_path_487, synth_cores=True).with_units("scalefree")
 
 
@@ -46,6 +46,7 @@ def test_filter_take(core_path_475, core_path_487):
 
 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)
+    print(ds.header.file.unit_convention)
     original_length = len(ds)
     original_redshift_range = ds.z_range
     output = tmp_path / "synth_gals.hdf5"