From 90938a2255f6529793a977762c70994640725128 Mon Sep 17 00:00:00 2001 From: Fridolin Glatter Date: Thu, 28 Nov 2024 12:32:50 +0100 Subject: [PATCH] Test adding data to scalar variable raises --- ixmp4/data/db/optimization/variable/repository.py | 13 ++++++++----- tests/core/test_optimization_variable.py | 8 ++++++++ 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/ixmp4/data/db/optimization/variable/repository.py b/ixmp4/data/db/optimization/variable/repository.py index 76aa13f3..c8bed49e 100644 --- a/ixmp4/data/db/optimization/variable/repository.py +++ b/ixmp4/data/db/optimization/variable/repository.py @@ -178,11 +178,14 @@ def add_data(self, variable_id: int, data: dict[str, Any] | pd.DataFrame) -> Non index_list = [column.name for column in variable.columns] existing_data = pd.DataFrame(variable.data) - if not existing_data.empty: - existing_data.set_index(index_list, inplace=True) - variable.data = ( - data.set_index(index_list).combine_first(existing_data).reset_index() - ).to_dict(orient="list") # type: ignore[assignment] + if index_list: + data = data.set_index(index_list) + if not existing_data.empty: + existing_data.set_index(index_list, inplace=True) + data = data.combine_first(existing_data) + if index_list: + data = data.reset_index() + variable.data = data.to_dict(orient="list") # type: ignore[assignment] self.session.commit() diff --git a/tests/core/test_optimization_variable.py b/tests/core/test_optimization_variable.py index 183e65f5..0dec1d45 100644 --- a/tests/core/test_optimization_variable.py +++ b/tests/core/test_optimization_variable.py @@ -273,6 +273,14 @@ def test_variable_add_data(self, platform: ixmp4.Platform) -> None: ) assert_unordered_equality(expected, pd.DataFrame(variable_4.data)) + # Test adding to scalar variable raises + with pytest.raises( + OptimizationDataValidationError, + match="Trying to add data to unknown Columns!", + ): + variable_5 = run.optimization.variables.create("Variable 5") + variable_5.add(data={"foo": ["bar"], "levels": [1], "marginals": [0]}) + def test_variable_remove_data(self, platform: ixmp4.Platform) -> None: run = platform.runs.create("Model", "Scenario") indexset = run.optimization.indexsets.create("Indexset")