diff --git a/powersimdata/tests/mock_analyze.py b/powersimdata/tests/mock_analyze.py index c00bd6d2d..5174edf4e 100644 --- a/powersimdata/tests/mock_analyze.py +++ b/powersimdata/tests/mock_analyze.py @@ -23,6 +23,39 @@ def _ensure_ts_index(df): class MockAnalyze: + """A mock of a powersimdata.scenario.analyze.Analyze object. + + :param dict grid_attrs: fields to be added to grid. + :param pandas.DataFrame congl: dummy congl + :param pandas.DataFrame congu: dummy congu + :param dict ct: dummy ct + :param pandas.DataFrame demand: dummy demand + :param pandas.DataFrame lmp: dummy lmp + :param pandas.DataFrame pf: dummy pf + :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 + """ + + exported_methods = [ + "get_congl", + "get_congu", + "get_ct", + "get_demand", + "get_grid", + "get_lmp", + "get_pf", + "gt_dcline_pf", + "get_pg", + "get_storage_pg", + "get_storage_e", + "get_solar", + "get_wind", + "get_hydro", + ] + def __init__( self, grid_attrs=None, @@ -34,26 +67,13 @@ def __init__( pf=None, pg=None, dcline_pf=None, + storage_e=None, storage_pg=None, solar=None, wind=None, hydro=None, ): - """Constructor. - - :param dict grid_attrs: fields to be added to grid. - :param pandas.DataFrame congl: dummy congl - :param pandas.DataFrame congu: dummy congu - :param dict ct: dummy ct - :param pandas.DataFrame demand: dummy demand - :param pandas.DataFrame lmp: dummy lmp - :param pandas.DataFrame pf: dummy pf - :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 - """ + """Constructor.""" self.grid = MockGrid(grid_attrs) self.congl = _ensure_ts_index(congl) self.congu = _ensure_ts_index(congu) @@ -63,6 +83,7 @@ def __init__( self.pf = _ensure_ts_index(pf) self.dcline_pf = _ensure_ts_index(dcline_pf) self.pg = _ensure_ts_index(pg) + self.storage_e = _ensure_ts_index(storage_e) self.storage_pg = _ensure_ts_index(storage_pg) self.solar = _ensure_ts_index(solar) self.wind = _ensure_ts_index(wind) @@ -123,6 +144,12 @@ def get_pg(self): """ return self.pg + def get_storage_e(self): + """Get storage E. + :return: (pandas.DataFrame) -- dummy storage_e + """ + return self.storage_e + def get_storage_pg(self): """Get storage PG. :return: (pandas.DataFrame) -- dummy storage_pg diff --git a/powersimdata/tests/mock_scenario.py b/powersimdata/tests/mock_scenario.py index 5012c0034..5491fb329 100644 --- a/powersimdata/tests/mock_scenario.py +++ b/powersimdata/tests/mock_scenario.py @@ -3,12 +3,16 @@ class MockScenario: - def __init__(self, grid_attrs=None, **kwargs): - """Constructor. + """ + :param dict grid_attrs: fields to be added to grid. + :param \\*\\*kwargs: collected keyword arguments to be passed to + MockAnalyze init. + """ + + _setattr_allowlist = {"state", "info"} - :param dict grid_attrs: fields to be added to grid. - :param pandas.DataFrame pg: dummy pg - """ + def __init__(self, grid_attrs=None, **kwargs): + """Constructor.""" self.state = MockAnalyze(grid_attrs, **kwargs) self.info = { "id": "111", @@ -34,3 +38,26 @@ def __init__(self, grid_attrs=None, **kwargs): def __class__(self): """If anyone asks, I'm a Scenario object!""" return Scenario + + def __dir__(self): + return sorted(super().__dir__() + list(self.state.exported_methods)) + + def __getattr__(self, name): + if name in self.state.exported_methods: + return getattr(self.state, name) + elif hasattr(self.state, "__getattr__"): + return self.state.__getattr__(name) + else: + raise AttributeError( + f"Scenario object in {self.state.name} state " + f"has no attribute {name}" + ) + + def __setattr__(self, name, value): + if name in self._setattr_allowlist: + super().__setattr__(name, value) + elif name in self.state.exported_methods: + raise AttributeError( + f"{name} is exported from Scenario.state, edit it there if necessary" + ) + super().__setattr__(name, value)