diff --git a/src/delta_barth/analysis/forecast.py b/src/delta_barth/analysis/forecast.py index 0a650f0..e15cf71 100644 --- a/src/delta_barth/analysis/forecast.py +++ b/src/delta_barth/analysis/forecast.py @@ -11,6 +11,7 @@ from xgboost import XGBRegressor from delta_barth.analysis import parse from delta_barth.api.requests import ( SalesPrognosisResponse, + SalesPrognosisResponseEntry, SalesPrognosisResults, SalesPrognosisResultsExport, get_sales_prognosis_data, @@ -50,6 +51,10 @@ def _parse_api_resp_to_df( """ data = resp.model_dump()["daten"] + if not data: + target_features = SalesPrognosisResponseEntry.__annotations__.keys() + data = {feat: [] for feat in target_features} + return pd.DataFrame(data) diff --git a/tests/analysis/test_forecast.py b/tests/analysis/test_forecast.py index d1c5666..485fdcf 100644 --- a/tests/analysis/test_forecast.py +++ b/tests/analysis/test_forecast.py @@ -8,6 +8,7 @@ from pydantic import ValidationError import delta_barth.analysis.forecast from delta_barth.analysis import forecast as fc +from delta_barth.api.requests import SalesPrognosisResponse, SalesPrognosisResponseEntry from delta_barth.errors import STATUS_HANDLER from delta_barth.types import DualDict, PipeResult @@ -102,16 +103,14 @@ def sales_data_real_preproc(sales_data_real, feature_map) -> pd.DataFrame: def test_parse_api_resp_to_df(exmpl_api_sales_prognosis_resp): resp = exmpl_api_sales_prognosis_resp df = fc._parse_api_resp_to_df(resp) - features = set( - ( - "artikelId", - "warengruppeId", - "firmaId", - "betrag", - "menge", - "buchungsDatum", - ) - ) + features = set(SalesPrognosisResponseEntry.__annotations__.keys()) + assert all(col in features for col in df.columns) + + +def test_parse_api_resp_to_df_empty(): + resp = SalesPrognosisResponse(daten=tuple()) + df = fc._parse_api_resp_to_df(resp) + features = set(SalesPrognosisResponseEntry.__annotations__.keys()) assert all(col in features for col in df.columns)