Fix tests for add_data

Author: glatterf42

ixmp4/data/db/base.py

@@ -488,9 +488,7 @@ def collect_indexsets_to_check(self) -> dict[str, list[Any]]:
         IndexSet.elements."""
         return {column.name: column.indexset.elements for column in self.columns}
 
-    @validates("data")
-    def validate_data(self, key, data: dict[str, Any]):
-        data_frame: pd.DataFrame = pd.DataFrame.from_dict(data)
+    def _validate_data(self, data_frame: pd.DataFrame, data: dict[str, Any]) -> None:
         # TODO for all of the following, we might want to create unique exceptions
         # Could we make both more specific by specifiying missing/extra columns?
         if len(data_frame.columns) < len(self.columns):
@@ -524,4 +522,8 @@ def validate_data(self, key, data: dict[str, Any]):
                 "and Columns it is constrained to!"
             )
 
+    @validates("data")
+    def validate_data(self, key, data: dict[str, Any]):
+        data_frame: pd.DataFrame = pd.DataFrame.from_dict(data)
+        self._validate_data(data_frame=data_frame, data=data)
         return data_frame.to_dict(orient="list")

ixmp4/data/db/optimization/parameter/model.py

@@ -1,29 +1,14 @@
-from typing import ClassVar
+from typing import Any, ClassVar
 
-from sqlalchemy import Column as sqlaColumn
-from sqlalchemy import Table
+import pandas as pd
+from sqlalchemy.orm import validates
 
 from ixmp4 import db
 from ixmp4.data import types
 from ixmp4.data.abstract import optimization as abstract
-from ixmp4.data.db.unit import Unit
 
 from .. import Column, base
 
-# Many Parameters can refer to many Units
-# note for a Core table, we use the sqlalchemy.Column construct,
-# not sqlalchemy.orm.mapped_column
-
-# TODO Is this enough/correct? This follows many-to-many currently with units:
-# But does that work?
-
-parameter_unit_association_table = Table(
-    "optimization_parameter_unit_association_table",
-    base.BaseModel.metadata,
-    sqlaColumn("parameter__id", db.ForeignKey("optimization_parameter.id")),
-    sqlaColumn("unit__id", db.ForeignKey("unit.id")),
-)
-
 
 class Parameter(base.BaseModel, base.OptimizationDataMixin, base.UniqueNameRunIDMixin):
     # NOTE: These might be mixin-able, but would require some abstraction
@@ -33,12 +18,13 @@ class Parameter(base.BaseModel, base.OptimizationDataMixin, base.UniqueNameRunID
 
     # constrained_to_indexsets: ClassVar[list[str] | None] = None
 
-    values: types.JsonList = db.Column(db.JsonType, nullable=False, default=[])
-    units: types.Mapped[list["Unit"]] = db.relationship(
-        secondary=parameter_unit_association_table
-    )
-    # TODO: need some kind of primaryjoin adaption and unit_ids so that each unit_id is
-    # foreignkeyed to Unit.id correctly
-
     # TODO Same as in table/model.py
     columns: types.Mapped[list["Column"]] = db.relationship()  # type: ignore
+
+    @validates("data")
+    def validate_data(self, key, data: dict[str, Any]):
+        data_frame: pd.DataFrame = pd.DataFrame.from_dict(data)
+        data_frame_to_validate = data_frame.drop(columns=["values", "units"])
+
+        self._validate_data(data_frame=data_frame_to_validate, data=data)
+        return data_frame.to_dict(orient="list")

ixmp4/data/db/optimization/parameter/repository.py

@@ -5,6 +5,7 @@
 from ixmp4 import db
 from ixmp4.data.abstract import optimization as abstract
 from ixmp4.data.auth.decorators import guard
+from ixmp4.data.db.unit import Unit
 
 from .. import ColumnRepository, base
 from .docs import ParameterDocsRepository
@@ -150,26 +151,24 @@ def add_data(self, parameter_id: int, data: dict[str, Any] | pd.DataFrame) -> No
             data = pd.DataFrame.from_dict(data=data)
         parameter = self.get_by_id(id=parameter_id)
 
-        try:
-            values = data.pop(item="values").to_list()
-        except KeyError as e:
-            raise KeyError("Parameter.data must include a 'values' column!") from e
+        missing_columns = set(["values", "units"]) - set(data.columns)
+        assert (
+            not missing_columns
+        ), f"Parameter.data must include the column(s): {' ,'.join(missing_columns)}!"
 
-        try:
-            units = [
+        # Can use a set for now, need full column if we care about order
+        for unit_name in set(data["units"]):
+            try:
                 self.backend.units.get(name=unit_name)
-                for unit_name in data.pop(item="units")
-            ]
-        except KeyError as e:
-            raise KeyError("Parameter.data must include a 'units' column!") from e
+            except Unit.NotFound as e:
+                # TODO Add a helpful hint on how to check defined Units
+                raise Unit.NotFound(
+                    message=f"'{unit_name}' is not defined for this Platform!"
+                ) from e
 
         parameter.data = pd.concat(
             [pd.DataFrame.from_dict(parameter.data), data]
         ).to_dict(orient="list")
-        parameter.values = parameter.values + values
-
-        # TODO does this actually work? Do we set the relationships correctly here?
-        parameter.units = parameter.units + units
 
         self.session.add(parameter)
         self.session.commit()

tests/data/test_optimization_parameter.py

@@ -49,8 +49,6 @@ def test_create_parameter(self, test_mp, request):
         assert parameter.run__id == run.id
         assert parameter.name == "Parameter"
         assert parameter.data == {}  # JsonDict type currently requires a dict, not None
-        assert parameter.values == []
-        assert parameter.units == []
         assert parameter.columns[0].name == "Indexset"
         assert parameter.columns[0].constrained_to_indexset == indexset_1.id
 
@@ -161,8 +159,6 @@ def test_parameter_add_data(self, test_mp, request):
         parameter = test_mp.backend.optimization.parameters.get(
             run_id=run.id, name="Parameter"
         )
-        assert parameter.values == test_data_1.pop("values")
-        assert [unit.name for unit in parameter.units] == test_data_1.pop("units")
         assert parameter.data == test_data_1
 
         parameter_2 = test_mp.backend.optimization.parameters.create(
@@ -171,7 +167,9 @@ def test_parameter_add_data(self, test_mp, request):
             constrained_to_indexsets=[indexset_1.name, indexset_2.name],
         )
 
-        with pytest.raises(KeyError, match="must include a 'values' column!"):
+        with pytest.raises(
+            AssertionError, match=r"must include the column\(s\): values!"
+        ):
             test_mp.backend.optimization.parameters.add_data(
                 parameter_id=parameter_2.id,
                 data=pd.DataFrame(
@@ -183,7 +181,9 @@ def test_parameter_add_data(self, test_mp, request):
                 ),
             )
 
-        with pytest.raises(KeyError, match="must include a 'units' column!"):
+        with pytest.raises(
+            AssertionError, match=r"must include the column\(s\): units!"
+        ):
             test_mp.backend.optimization.parameters.add_data(
                 parameter_id=parameter_2.id,
                 data=pd.DataFrame(
@@ -232,8 +232,6 @@ def test_parameter_add_data(self, test_mp, request):
         parameter_2 = test_mp.backend.optimization.parameters.get(
             run_id=run.id, name="Parameter 2"
         )
-        assert parameter_2.values == test_data_2.pop("values")
-        assert [unit.name for unit in parameter_2.units] == test_data_2.pop("units")
         assert parameter_2.data == test_data_2
 
         # Test order is conserved with varying types and upon later addition of data
@@ -258,17 +256,8 @@ def test_parameter_add_data(self, test_mp, request):
         parameter_3 = test_mp.backend.optimization.parameters.get(
             run_id=run.id, name="Parameter 3"
         )
-        assert parameter_3.values == test_data_3.pop("values")
-        assert [unit.name for unit in parameter_3.units] == test_data_3.pop("units")
         assert parameter_3.data == test_data_3
 
-        # Repopulate test_data after pop()
-        test_data_3 = {
-            "Column 1": ["bar", "foo", ""],
-            "Column 2": [2, 3, 1],
-            "values": ["3", 2.0, 1],
-            "units": [unit_3.name, unit_2.name, unit.name],
-        }
         test_data_4 = {
             "Column 1": ["foo", "", "bar"],
             "Column 2": [2, 3, 1],
@@ -281,15 +270,10 @@ def test_parameter_add_data(self, test_mp, request):
         parameter_3 = test_mp.backend.optimization.parameters.get(
             run_id=run.id, name="Parameter 3"
         )
-        assert parameter_3.values == test_data_3.pop("values") + test_data_4.pop(
-            "values"
-        )
-        assert [unit.name for unit in parameter_3.units] == test_data_3.pop(
-            "units"
-        ) + test_data_4.pop("units")
-        assert parameter_3.data == pd.DataFrame([test_data_3, test_data_4]).to_dict(
-            orient="list"
-        )
+        test_data_5 = test_data_3.copy()
+        for key, value in test_data_4.items():
+            test_data_5[key].extend(value)
+        assert parameter_3.data == test_data_5
 
     # def test_list_parameter(self, test_mp, request):
     #     test_mp = request.getfixturevalue(test_mp)