wip

powersimdata/input/tests/test_transform_profile.py

@@ -1,14 +1,14 @@
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 import numpy as np
-import pandas as pd
 import pytest
 from numpy.testing import assert_almost_equal
 
 from powersimdata.input.change_table import ChangeTable
 from powersimdata.input.grid import Grid
 from powersimdata.input.transform_grid import TransformGrid
 from powersimdata.input.transform_profile import TransformProfile
+from powersimdata.tests.mock_input_data import MockInputData
 
 interconnect = ["Western"]
 param = {
@@ -34,11 +34,6 @@ def get_plant_with_resource(base_grid, resource):
     return plant_id
 
 
-def _get_plant_ids_for_type(grid, type):
-    plant_ids = list(grid.plant.query("type in @profile_type[@type]").index)
-    return plant_ids
-
-
 def get_change_table_for_zone_scaling(base_grid, resource):
     n_zone = param["n_zone_to_scale"]
     zones = get_zone_with_resource(base_grid, resource)
@@ -140,7 +135,9 @@ def _check_plants_are_scaled(ct, base_grid, raw_profile, resource):
 
 def _check_new_plants_are_added(ct, base_grid, raw_profile, resource):
     n_plant = param["n_plant_to_add"]
-    plant_id_type = _get_plant_ids_for_type(base_grid, resource)
+    plant_id_type = (
+        base_grid.plant.isin(profile_type[resource]).query("type == True").index
+    )
     base_profile = (
         raw_profile[plant_id_type] * base_grid.plant.loc[plant_id_type, "Pmax"]
     )
@@ -188,76 +185,45 @@ def _check_profile_of_new_plants_are_produced_correctly(
         assert new_profile[c].equals(neighbor_profile.multiply(new_plant_pmax[i]))
 
 
-@pytest.fixture(scope="module")
-def random():
-    return np.random.default_rng(seed=6669)
-
-
 @pytest.fixture(scope="module")
 def base_grid():
     grid = Grid(interconnect)
     return grid
 
 
-def _create_fake_profile(random, columns):
-    start_time = "2016-01-01 00:00:00"
-    times = pd.date_range(start_time, periods=24, freq="H")
-    index = pd.Index(times, name="UTC Time")
-    data = random.uniform(low=0, high=1, size=(len(times), len(columns)))
-    fake_profile = pd.DataFrame(data=data, index=index, columns=columns)
-    return fake_profile
+@pytest.fixture(scope="module")
+def input_data(base_grid):
+    mock_input_data = MockInputData(base_grid)
+    return mock_input_data
 
 
-@pytest.fixture(scope="module")
-def raw_hydro(base_grid, random):
-    plant_ids = _get_plant_ids_for_type(base_grid, "hydro")
-    fake_hydro_profile = _create_fake_profile(random, plant_ids)
-    return fake_hydro_profile
+@pytest.fixture(scope="module", autouse=True)
+def mock_input_data_class(input_data):
+    with patch(
+        "powersimdata.input.transform_profile.InputData"
+    ) as mock_input_data_class:
+        mock_input_data_class.return_value = input_data
+        yield
 
 
 @pytest.fixture(scope="module")
-def raw_wind(base_grid, random):
-    plant_ids = _get_plant_ids_for_type(base_grid, "wind")
-    fake_wind_profile = _create_fake_profile(random, plant_ids)
-    return fake_wind_profile
+def raw_hydro(input_data):
+    return input_data.get_data({}, "hydro")
 
 
 @pytest.fixture(scope="module")
-def raw_solar(base_grid, random):
-    plant_ids = _get_plant_ids_for_type(base_grid, "solar")
-    fake_solar_profile = _create_fake_profile(random, plant_ids)
-    return fake_solar_profile
+def raw_wind(input_data):
+    return input_data.get_data({}, "wind")
 
 
 @pytest.fixture(scope="module")
-def raw_demand(base_grid, random):
-    zone_ids = list(base_grid.id2zone.keys())
-    fake_demand_profile = _create_fake_profile(random, zone_ids)
-    return fake_demand_profile
+def raw_solar(input_data):
+    return input_data.get_data({}, "solar")
 
 
-@pytest.fixture(scope="module", autouse=True)
-def mock_input_data(raw_demand, raw_hydro, raw_wind, raw_solar):
-    with patch(
-        "powersimdata.input.transform_profile.InputData"
-    ) as mock_input_data_class:
-
-        def mock_get_data(scenario_info, field_name):
-            if field_name == "demand":
-                return raw_demand
-            elif field_name == "hydro":
-                return raw_hydro
-            elif field_name == "wind":
-                return raw_wind
-            elif field_name == "solar":
-                return raw_solar
-
-            raise ValueError(f"No return specified for {field_name}!")
-
-        mock_input_data_object = MagicMock()
-        mock_input_data_object.get_data.side_effect = mock_get_data
-        mock_input_data_class.return_value = mock_input_data_object
-        yield
+@pytest.fixture(scope="module")
+def raw_demand(input_data):
+    return input_data.get_data({}, "demand")
 
 
 def test_demand_is_scaled(base_grid, raw_demand):

powersimdata/tests/mock_input_data.py

@@ -0,0 +1,63 @@
+from functools import lru_cache
+
+import numpy as np
+import pandas as pd
+
+from powersimdata.input.grid import Grid
+
+
+def _get_plant_ids_for_type(grid, type):
+    profile_type = {
+        "wind": {"wind", "wind_offshore"},
+        "solar": {"solar"},
+        "hydro": {"hydro"},
+        "demand": {"demand"},
+    }
+
+    plant_ids = list(grid.plant.query("type in @profile_type[@type]").index)
+    return plant_ids
+
+
+class MockInputData:
+    def __init__(
+        self,
+        grid: Grid,
+        start_time="2016-01-01 00:00:00",
+        periods=24,
+        freq="H",
+        random_seed=6669,
+    ):
+        self._grid = grid
+        self._start_time = start_time
+        self._periods = periods
+        self._freq = freq
+        self._random = np.random.default_rng(seed=random_seed)
+
+    def get_data(self, scenario_info, field_name):
+        if field_name == "demand":
+            return self._get_demand()
+        elif field_name in ["hydro", "wind", "solar"]:
+            return self._get_resource_profile(field_name)
+
+        raise ValueError(f"No return specified for {field_name}!")
+
+    @lru_cache
+    def _get_demand(self):
+        zone_ids = list(self._grid.id2zone.keys())
+        fake_demand_profile = self._create_fake_profile(zone_ids)
+        return fake_demand_profile
+
+    @lru_cache
+    def _get_resource_profile(self, resource_type):
+        plant_ids = _get_plant_ids_for_type(self._grid, resource_type)
+        fake_resource_profile = self._create_fake_profile(plant_ids)
+        return fake_resource_profile
+
+    def _create_fake_profile(self, columns):
+        times = pd.date_range(
+            start=self._start_time, periods=self._periods, freq=self._freq
+        )
+        index = pd.Index(times, name="UTC Time")
+        data = self._random.uniform(low=0, high=1, size=(len(times), len(columns)))
+        fake_profile = pd.DataFrame(data=data, index=index, columns=columns)
+        return fake_profile