use consumption_price_sensor from flex_context

Signed-off-by: Victor Garcia Reolid <[email protected]>
FlexMeasures · Jul 25, 2023 · 98d4a16 · 98d4a16
1 parent fbffd2e
commit 98d4a16
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 6 deletions.
diff --git a/flexmeasures/data/models/planning/shiftable_load.py b/flexmeasures/data/models/planning/shiftable_load.py
@@ -61,7 +61,7 @@ def compute(self) -> pd.Series | None:
         belief_time = self.belief_time
         sensor = self.sensor
 
-        consumption_price_sensor: Sensor = self.flex_model.get(
+        consumption_price_sensor: Sensor = self.flex_context.get(
             "consumption_price_sensor"
         )
         duration: timedelta = self.flex_model.get("duration")

diff --git a/flexmeasures/data/models/planning/tests/test_shiftable_load.py b/flexmeasures/data/models/planning/tests/test_shiftable_load.py
@@ -28,18 +28,22 @@ def test_shiftable_scheduler(
     epex_da = Sensor.query.filter(Sensor.name == "epex_da").one_or_none()
 
     flex_model = {
-        "consumption-price-sensor": epex_da.id,
         "duration": "PT4H",
         "load-type": load_type,
         "power": 4,
     }
 
+    flex_context = {
+        "consumption-price-sensor": epex_da.id,
+    }
+
     scheduler = ShiftableLoadScheduler(
         shiftable_load,
         start,
         end,
         resolution,
         flex_model=flex_model,
+        flex_context=flex_context,
     )
     schedule = scheduler.compute()
 
@@ -68,18 +72,22 @@ def test_duration_exceeds_planning_window(
     epex_da = Sensor.query.filter(Sensor.name == "epex_da").one_or_none()
 
     flex_model = {
-        "consumption-price-sensor": epex_da.id,
         "duration": "PT48H",
         "load-type": load_type,
         "power": 4,
     }
 
+    flex_context = {
+        "consumption-price-sensor": epex_da.id,
+    }
+
     scheduler = ShiftableLoadScheduler(
         shiftable_load,
         start,
         end,
         resolution,
         flex_model=flex_model,
+        flex_context=flex_context,
     )
     schedule = scheduler.compute()
 
@@ -100,21 +108,24 @@ def test_shiftable_scheduler_time_restrictions(add_battery_assets, shiftable_loa
     # time parameters
 
     flex_model = {
-        "consumption-price-sensor": epex_da.id,
         "duration": "PT4H",
         "load-type": "SHIFTABLE",
         "power": 4,
         "time-restrictions": [
             {"start": "2015-01-02T08:00:00+01:00", "duration": "PT2H"}
         ],
     }
+    flex_context = {
+        "consumption-price-sensor": epex_da.id,
+    }
 
     scheduler = ShiftableLoadScheduler(
         shiftable_load,
         start,
         end,
         resolution,
         flex_model=flex_model,
+        flex_context=flex_context,
     )
     schedule = scheduler.compute()
 
@@ -148,7 +159,6 @@ def test_breakable_scheduler_time_restrictions(add_battery_assets, shiftable_loa
     # time parameters
 
     flex_model = {
-        "consumption-price-sensor": epex_da.id,
         "duration": "PT4H",
         "load-type": "BREAKABLE",
         "power": 4,
@@ -160,12 +170,17 @@ def test_breakable_scheduler_time_restrictions(add_battery_assets, shiftable_loa
         ],
     }
 
+    flex_context = {
+        "consumption-price-sensor": epex_da.id,
+    }
+
     scheduler = ShiftableLoadScheduler(
         shiftable_load,
         start,
         end,
         resolution,
         flex_model=flex_model,
+        flex_context=flex_context,
     )
     schedule = scheduler.compute()
 
@@ -200,19 +215,22 @@ def test_impossible_schedules(
     epex_da = Sensor.query.filter(Sensor.name == "epex_da").one_or_none()
 
     flex_model = {
-        "consumption-price-sensor": epex_da.id,
         "duration": "PT4H",
         "load-type": load_type,
         "power": 4,
         "time-restrictions": time_restrictions,
     }
+    flex_context = {
+        "consumption-price-sensor": epex_da.id,
+    }
 
     scheduler = ShiftableLoadScheduler(
         shiftable_load,
         start,
         end,
         resolution,
         flex_model=flex_model,
+        flex_context=flex_context,
     )
 
     with pytest.raises(ValueError):