diff --git a/tests/test_adjuster.py b/tests/test_adjuster.py index 761e731..fe80278 100644 --- a/tests/test_adjuster.py +++ b/tests/test_adjuster.py @@ -4,6 +4,8 @@ import pandas as pd import pytest +from pvsite_datamodel.sqlmodels import SiteAssetType + from india_forecast_app.adjuster import ( adjust_forecast_with_adjuster, get_me_values, @@ -80,6 +82,9 @@ def test_adjust_forecast_with_adjuster(db_session, sites, generation_db_values, { "forecast_power_kw": [1, 2, 3, 4, 5], "horizon_minutes": [15, 30, 45, 60, 1200], + "start_utc": [ + pd.Timestamp("2024-11-01 03:00:00") + pd.Timedelta(f"{i}H") for i in range(0, 5) + ], } ) @@ -100,6 +105,9 @@ def test_adjust_forecast_with_adjuster_no_values(db_session, sites): { "forecast_power_kw": [1, 2, 3, 4, 5], "horizon_minutes": [15, 30, 45, 60, 1200], + "start_utc": [ + pd.Timestamp("2024-11-01 03:00:00") + pd.Timedelta(f"{i}H") for i in range(0, 5) + ], } ) @@ -111,7 +119,7 @@ def test_adjust_forecast_with_adjuster_no_values(db_session, sites): assert forecast_values_df["forecast_power_kw"].sum() == 15 -@pytest.mark.parametrize("asset_type", ["pv", "wind"]) +@pytest.mark.parametrize("asset_type", [SiteAssetType.pv, SiteAssetType.wind]) def test_zero_out_night_time_for_pv(asset_type, db_session, sites): """ Test for zero_out_nighttime """ forecast_values_df = pd.DataFrame( @@ -132,7 +140,7 @@ def test_zero_out_night_time_for_pv(asset_type, db_session, sites): assert len(forecast_values_df) == 5 night_sum = forecast_values_df["forecast_power_kw"][0:2].sum() - if asset_type == "pv": + if asset_type == SiteAssetType.pv: assert night_sum == 0 else: assert night_sum > 0