add horizon to csv route & format func

src/india_api/internal/inputs/indiadb/test_csv.py

@@ -0,0 +1,71 @@
+import logging
+
+import pandas as pd
+from fastapi import HTTPException
+import pytest
+
+from india_api.internal import PredictedPower, ActualPower, SiteProperties
+
+from pvsite_datamodel.sqlmodels import APIRequestSQL
+
+from .client import Client
+from .conftest import forecast_values
+from ...models import ForecastHorizon
+from ...service.csv import format_csv_and_created_time
+
+log = logging.getLogger(__name__)
+
+# TODO add list of test that are here
+
+
+@pytest.fixture()
+def client(engine, db_session):
+    """Hooks Client into pytest db_session fixture"""
+    client = Client(database_url=str(engine.url))
+    client.session = db_session
+
+    return client
+
+# Skip for now
+@pytest.mark.skip(reason="Not finished yet")
+class TestCsvExport:
+    def test_format_csv_and_created_time(self, client, forecast_values_wind) -> None:
+        """Test the format_csv_and_created_time function."""
+        forecast_values_wind = client.get_predicted_wind_power_production_for_location(
+            location="testID"
+        )
+        assert forecast_values_wind is not None
+        assert len(forecast_values_wind) > 0
+        assert isinstance(forecast_values_wind[0], PredictedPower)
+
+        result = format_csv_and_created_time(
+            forecast_values_wind,
+            ForecastHorizon.latest,
+        )
+        assert isinstance(result, tuple)
+        assert isinstance(result[0], pd.DataFrame)
+        assert isinstance(result[1], pd.Timestamp)
+        logging.info(f"CSV created at: {result[1]}")
+        logging.info(f"CSV content: {result[0].head()}")
+        # Check the shape of the DataFrame
+        # The shape should match the number of forecast values
+        # and the number of columns in the DataFrame
+        # The DataFrame should have 3 columns: Date [IST], Time, PowerMW
+        assert result[0].shape[1] == 3
+        # Check the first row of the DataFrame
+        # The date of the first row should be the nearest rounded 15min from now
+        rounded_15_min = pd.Timestamp.now(tz="Asia/Kolkata").round("15min")
+        assert result[0].iloc[0]["Time"] == rounded_15_min.strftime("%H:%M")
+        # Check the number of rows in the DataFrame
+        # For the latest forecast, it should be the number of
+        # forecast values after now
+        forecast_values_from_now = [
+            value for value in forecast_values_wind if value.Time >= rounded_15_min
+        ]
+        assert result[0].shape[0] == len(forecast_values_from_now)
+        # Check the column names
+        assert list(result[0].columns) == ["Date [IST]", "Time", "PowerMW"]
+        # Check the data types of the columns
+        assert result[0]["Date [IST]"].dtype == "datetime64[ns, Asia/Kolkata]"
+        assert result[0]["Time"].dtype == "object"
+        assert result[0]["PowerMW"].dtype == "float64"

src/india_api/internal/service/csv.py

@@ -2,9 +2,10 @@
 from datetime import datetime
 
 from india_api.internal import PredictedPower
+from india_api.internal.models import ForecastHorizon
 
 
-def format_csv_and_created_time(values: list[PredictedPower]) -> (pd.DataFrame, datetime):
+def format_csv_and_created_time(values: list[PredictedPower], forecast_horizon: ForecastHorizon) -> (pd.DataFrame, datetime):
     """
     Format the predicted power values into a pandas dataframe ready for CSV export.
 
@@ -26,16 +27,20 @@ def format_csv_and_created_time(values: list[PredictedPower]) -> (pd.DataFrame,
     df["Date [IST]"] = df["Time"].dt.date
     # create start and end time column and only show HH:MM
     df["Start Time [IST]"] = df["Time"].dt.strftime("%H:%M")
-    df["End Time [IST]"] = (df["Time"] + pd.to_timedelta("15T")).dt.strftime("%H:%M")
+    df["End Time [IST]"] = (df["Time"] + pd.to_timedelta("15min")).dt.strftime("%H:%M")
+
+    now_ist = pd.Timestamp.now(tz="Asia/Kolkata")
+    if forecast_horizon == ForecastHorizon.day_ahead:
+        # only get tomorrow's results, for IST time.
+        tomorrow = now_ist + pd.Timedelta(days=1)
+        df = df[df["Date [IST]"] == tomorrow.date()]
+    elif forecast_horizon == ForecastHorizon.latest:
+        # only get results from now onwards, for IST time.
+        df = df[df["Time"] >= now_ist]
 
     # combine start and end times
     df["Time"] = df["Start Time [IST]"].astype(str) + " - " + df["End Time [IST]"].astype(str)
 
-    # only get tomorrows results. This is for IST time.
-    now_ist = pd.Timestamp.now(tz="Asia/Kolkata")
-    tomorrow = now_ist + pd.Timedelta(days=1)
-    df = df[df["Date [IST]"] == tomorrow.date()]
-
     # get the max created time
     created_time = df["CreatedTime"].max()
 

src/india_api/internal/service/regions.py

@@ -186,35 +186,57 @@ def get_forecast_timeseries_route(
     response_class=FileResponse,
     include_in_schema=False,
 )
-def get_forecast_da_csv(
+def get_forecast_csv(
     source: ValidSourceDependency,
     region: str,
     db: DBClientDependency,
     auth: dict = Depends(auth),
+    forecast_horizon: Optional[ForecastHorizon] = ForecastHorizon.latest,
 ):
     """
     Route to get the day ahead forecast as a CSV file.
+    By default, the CSV file will be for the latest forecast, from now forwards.
+    The forecast_horizon can be set to 'latest' or 'day_ahead'.
+    - latest: The latest forecast, from now forwards.
+    - day_ahead: The forecast for the next day, from 00:00.
     """
 
-    forcasts: GetForecastGenerationResponse = get_forecast_timeseries_route(
+    if forecast_horizon is not None:
+        if forecast_horizon not in [ForecastHorizon.latest, ForecastHorizon.day_ahead]:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail=f"Invalid forecast_horizon {forecast_horizon}. Must be 'latest' or 'day_ahead'.",
+            )
+
+    forecasts: GetForecastGenerationResponse = get_forecast_timeseries_route(
         source=source,
         region=region,
         db=db,
         auth=auth,
-        forecast_horizon=ForecastHorizon.day_ahead,
+        forecast_horizon=forecast_horizon,
         smooth_flag=False,
     )
 
     # format to dataframe
-    df, created_time = format_csv_and_created_time(forcasts.values)
+    df, created_time = format_csv_and_created_time(forecasts.values, forecast_horizon=forecast_horizon)
 
     # make file format
     now_ist = pd.Timestamp.now(tz="Asia/Kolkata")
     tomorrow_ist = df["Date [IST]"].iloc[0]
-    csv_file_path = f"{region}_{source}_da_{tomorrow_ist}.csv"
+    match forecast_horizon:
+        case ForecastHorizon.latest:
+            forecast_type = "intraday"
+        case ForecastHorizon.day_ahead:
+            forecast_type = "da"
+        case _:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail=f"Invalid forecast_horizon {forecast_horizon}. Must be 'latest' or 'day_ahead'.",
+            )
+    csv_file_path = f"{region}_{source}_{forecast_type}_{tomorrow_ist}.csv"
 
     description = (
-        f"Forecast for {region} for {source} for {tomorrow_ist}. "
+        f"Forecast for {region} for {source}, {forecast_type}, for {tomorrow_ist}. "
         f"The Forecast was created at {created_time} and downloaded at {now_ist}"
     )