feat: add storage_pg entry to MockAnalyze

Breakthrough-Energy · Jan 28, 2021 · c7967eb · c7967eb
1 parent a84b126
commit c7967eb
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/powersimdata/tests/mock_analyze.py b/powersimdata/tests/mock_analyze.py
@@ -32,6 +32,7 @@ def __init__(
         demand=None,
         lmp=None,
         pg=None,
+        storage_pg=None,
         solar=None,
         wind=None,
         hydro=None,
@@ -45,6 +46,7 @@ def __init__(
         :param pandas.DataFrame demand: dummy demand
         :param pandas.DataFrame lmp: dummy lmp
         :param pandas.DataFrame pg: dummy pg
+        :param pandas.DataFrame storage_pg: dummy storage_pg
         :param pandas.DataFrame solar: dummy solar
         :param pandas.DataFrame wind: dummy wind
         :param pandas.DataFrame hydro: dummy hydro
@@ -56,6 +58,7 @@ def __init__(
         self.demand = _ensure_ts_index(demand)
         self.lmp = _ensure_ts_index(lmp)
         self.pg = _ensure_ts_index(pg)
+        self.storage_pg = _ensure_ts_index(storage_pg)
         self.solar = _ensure_ts_index(solar)
         self.wind = _ensure_ts_index(wind)
         self.hydro = _ensure_ts_index(hydro)
@@ -103,6 +106,12 @@ def get_pg(self):
         """
         return self.pg
 
+    def get_storage_pg(self):
+        """Get storage PG.
+        :return: (pandas.DataFrame) -- dummy storage_pg
+        """
+        return self.storage_pg
+
     def get_solar(self):
         """Get solar.
         :return: (pandas.DataFrame) -- dummy solar