From 9471fb54b327cc9cce4a0115099e3bdf85aadeda Mon Sep 17 00:00:00 2001
From: Ian Menezes <ian.menezes@rte-france.com>
Date: Fri, 29 Nov 2024 17:05:09 +0100
Subject: [PATCH] Fixed unit test errors

---
 src/andromede/simulation/optimization.py    | 13 +++++++++----
 src/andromede/study/data.py                 | 14 ++++++++------
 tests/functional/test_investment_pathway.py |  2 +-
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/src/andromede/simulation/optimization.py b/src/andromede/simulation/optimization.py
index 3f5463c3..d411106f 100644
--- a/src/andromede/simulation/optimization.py
+++ b/src/andromede/simulation/optimization.py
@@ -623,8 +623,12 @@ def _solver_variable_name(
             if (self.context.full_var_name and self.context.tree_node)
             else ""
         )
-        scenario_suffix = f"_s{s}" if s is not None else ""
-        block_suffix = f"_t{t}" if t is not None else ""
+        scenario_suffix = (
+            f"_s{s}" if (s is not None and self.context.scenarios > 1) else ""
+        )
+        block_suffix = (
+            f"_t{t}" if (t is not None and self.context.block_length() > 1) else ""
+        )
 
         # Set solver var name
         # Externally, for the Solver, this variable will have a full name
@@ -653,10 +657,11 @@ def _create_variables(self) -> None:
                     )
 
                 time_indices: Iterable[Optional[int]] = [None]
-                if var_indexing.time:
+                if var_indexing.is_time_varying():
                     time_indices = self.context.get_time_indices(var_indexing)
+
                 scenario_indices: Iterable[Optional[int]] = [None]
-                if var_indexing.scenario:
+                if var_indexing.is_scenario_varying():
                     scenario_indices = self.context.get_scenario_indices(var_indexing)
 
                 for t, s in itertools.product(time_indices, scenario_indices):
diff --git a/src/andromede/study/data.py b/src/andromede/study/data.py
index c3a1b197..99df1640 100644
--- a/src/andromede/study/data.py
+++ b/src/andromede/study/data.py
@@ -191,7 +191,7 @@ def check_requirement(self, time: bool, scenario: bool) -> bool:
 
 
 @dataclass(frozen=True)
-class DatabaseIndex:
+class ComponentParameterIndex:
     component_id: str
     parameter_name: str
 
@@ -204,20 +204,22 @@ class DataBase:
     Data can have different structure : constant, varying in time or scenarios.
     """
 
-    _data: Dict[DatabaseIndex, AbstractDataStructure]
+    _data: Dict[ComponentParameterIndex, AbstractDataStructure]
 
     def __init__(self) -> None:
-        self._data: Dict[DatabaseIndex, AbstractDataStructure] = {}
+        self._data: Dict[ComponentParameterIndex, AbstractDataStructure] = {}
 
     def get_data(self, component_id: str, parameter_name: str) -> AbstractDataStructure:
-        return self._data[DatabaseIndex(component_id, parameter_name)]
+        return self._data[ComponentParameterIndex(component_id, parameter_name)]
 
     def add_data(
         self, component_id: str, parameter_name: str, data: AbstractDataStructure
     ) -> None:
-        self._data[DatabaseIndex(component_id, parameter_name)] = data
+        self._data[ComponentParameterIndex(component_id, parameter_name)] = data
 
-    def get_value(self, index: DatabaseIndex, timestep: int, scenario: int) -> float:
+    def get_value(
+        self, index: ComponentParameterIndex, timestep: int, scenario: int
+    ) -> float:
         if index in self._data:
             return self._data[index].get_value(timestep, scenario)
         else:
diff --git a/tests/functional/test_investment_pathway.py b/tests/functional/test_investment_pathway.py
index 81895ac0..38ef6e2b 100644
--- a/tests/functional/test_investment_pathway.py
+++ b/tests/functional/test_investment_pathway.py
@@ -95,7 +95,7 @@ def candidate() -> Component:
                 )
             ],
             objective_operational_contribution=(param("op_cost") * var("generation"))
-            .sum()
+            .time_sum()
             .expec(),
             objective_investment_contribution=param("invest_cost")
             * var("delta_invest"),