diff --git a/neso_solar_consumer/fetch_data.py b/neso_solar_consumer/fetch_data.py index 0267af8..01c21d8 100644 --- a/neso_solar_consumer/fetch_data.py +++ b/neso_solar_consumer/fetch_data.py @@ -8,9 +8,7 @@ import pandas as pd -def fetch_data( - resource_id: str, limit: int, columns: list, rename_columns: dict -) -> pd.DataFrame: +def fetch_data(resource_id: str, limit: int, columns: list, rename_columns: dict) -> pd.DataFrame: """ Fetch data from the NESO API and process it into a Pandas DataFrame. @@ -55,9 +53,7 @@ def fetch_data( return pd.DataFrame() -def fetch_data_using_sql( - sql_query: str, columns: list, rename_columns: dict -) -> pd.DataFrame: +def fetch_data_using_sql(sql_query: str, columns: list, rename_columns: dict) -> pd.DataFrame: """ Fetch data from the NESO API using an SQL query, process it, and return specific columns with renamed headers. @@ -98,4 +94,4 @@ def fetch_data_using_sql( except Exception as e: print(f"An error occurred: {e}") - return pd.DataFrame() \ No newline at end of file + return pd.DataFrame() diff --git a/neso_solar_consumer/format_forecast.py b/neso_solar_consumer/format_forecast.py index 42129bf..69fd29d 100644 --- a/neso_solar_consumer/format_forecast.py +++ b/neso_solar_consumer/format_forecast.py @@ -35,7 +35,7 @@ def format_to_forecast_sql(data: pd.DataFrame, model_tag: str, model_version: st forecast_values = [ ForecastValue( target_time=target_time, - expected_power_generation_megawatts=row["solar_forecast_kw"] + expected_power_generation_megawatts=row["solar_forecast_kw"], ).to_orm() ] diff --git a/neso_solar_consumer/save_forecasts.py b/neso_solar_consumer/save_forecasts.py index 6dd9183..abc5966 100644 --- a/neso_solar_consumer/save_forecasts.py +++ b/neso_solar_consumer/save_forecasts.py @@ -2,6 +2,7 @@ from nowcasting_datamodel.save.save import save + def save_forecasts_to_db(forecasts: list, session): """ Save a list of ForecastSQL objects to the database using the nowcasting_datamodel `save` function. diff --git a/tests/test_fetch_data.py b/tests/test_fetch_data.py index dc7ae79..b020a0b 100644 --- a/tests/test_fetch_data.py +++ b/tests/test_fetch_data.py @@ -26,7 +26,7 @@ """ -import pytest +import pytest # noqa: F401 from neso_solar_consumer.fetch_data import fetch_data, fetch_data_using_sql @@ -69,6 +69,4 @@ def test_data_consistency(): """ df_api = fetch_data(resource_id, limit, columns, rename_columns) df_sql = fetch_data_using_sql(sql_query, columns, rename_columns) - assert df_api.equals( - df_sql - ), "Data from fetch_data and fetch_data_using_sql are inconsistent!" + assert df_api.equals(df_sql), "Data from fetch_data and fetch_data_using_sql are inconsistent!" diff --git a/tests/test_save_forecasts.py b/tests/test_save_forecasts.py index ccd6867..6419b15 100644 --- a/tests/test_save_forecasts.py +++ b/tests/test_save_forecasts.py @@ -26,6 +26,7 @@ MODEL_NAME = "real_data_model" MODEL_VERSION = "1.0" + @pytest.fixture def db_session(): """ @@ -44,6 +45,7 @@ def db_session(): session.close() engine.dispose() + def test_save_real_forecasts(db_session): """ Integration test: Fetch real data, format it into forecasts, and save to the database.