From 6309035a89d7bdb0d18d6789781adcb77d07060b Mon Sep 17 00:00:00 2001 From: Khusham Bansal Date: Fri, 26 Dec 2025 04:35:16 +0530 Subject: [PATCH 1/2] Add pytest for csv file download --- .../backends/quartzdb/test_csv_download.py | 148 ++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 src/quartz_api/internal/backends/quartzdb/test_csv_download.py diff --git a/src/quartz_api/internal/backends/quartzdb/test_csv_download.py b/src/quartz_api/internal/backends/quartzdb/test_csv_download.py new file mode 100644 index 00000000..2bc6f94c --- /dev/null +++ b/src/quartz_api/internal/backends/quartzdb/test_csv_download.py @@ -0,0 +1,148 @@ +"""Test CSV download endpoint.""" + +import io +import pathlib + +import pandas as pd +import pytest +from fastapi.testclient import TestClient +from pyhocon import ConfigFactory +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session + +from quartz_api.cmd.main import _create_server +from quartz_api.internal import models + +from .client import Client +import datetime as dt +from pvsite_datamodel.read.model import get_or_create_model +from pvsite_datamodel.sqlmodels import ForecastSQL, ForecastValueSQL, LocationSQL + + +@pytest.fixture() +def client(engine: Engine, db_session: Session) -> Client: + """Hooks Client into pytest db_session fixture.""" + client = Client(database_url=str(engine.url)) + client.session = db_session + + return client + + +@pytest.fixture() +def test_client(client: Client) -> TestClient: + """Create a FastAPI test client.""" + # Load config and override routers to include 'regions' + conf_path = pathlib.Path(__file__).parent.parent.parent.parent / "cmd" / "server.conf" + conf = ConfigFactory.parse_file(conf_path.as_posix()) + + # Override config to include regions router + override_conf = ConfigFactory.from_dict({"api": {"routers": "regions"}}) + conf = override_conf.with_fallback(conf) + + # Create the FastAPI app + app = _create_server(conf=conf) + + # Override the database dependency with our test client + app.dependency_overrides[models.get_db_client] = lambda: client + + return TestClient(app) + + +class TestCsvDownload: + """Test CSV download functionality.""" + + def test_csv_download_success( + self, + test_client: TestClient, + client: Client, + forecast_values_wind: None, + ) -> None: + """Test successful CSV download from the API endpoint.""" + # Make request to CSV download endpoint + response = test_client.get( + "/wind/testID/forecast/csv", + params={"forecast_horizon": "latest"}, + ) + + # Assert successful response + assert response.status_code == 200 + assert response.headers["content-type"] == "text/csv; charset=utf-8" + assert "Content-Disposition" in response.headers + assert "attachment" in response.headers["Content-Disposition"] + assert ".csv" in response.headers["Content-Disposition"] + + # Parse CSV content + csv_content = response.text + df = pd.read_csv(io.StringIO(csv_content)) + + # Validate CSV structure + assert "Date [IST]" in df.columns + assert "Time" in df.columns + assert "PowerMW" in df.columns + assert len(df) > 0 + + def test_csv_download_day_ahead( + self, + test_client: TestClient, + client: Client, + forecast_values_wind: None, + ) -> None: + """Test CSV download with day_ahead forecast horizon.""" + + # 1. Setup: Create 'Day Ahead' forecast data (for tomorrow) + # Get the wind site + site = client.session.query(LocationSQL).filter(LocationSQL.asset_type == "wind").first() + + # Create a forecast timestamped now + now = dt.datetime.now(tz=dt.UTC) + + # Create a model entry + ml_model = get_or_create_model(client.session, "windnet_india_adjust") + + forecast = ForecastSQL( + location_uuid=site.location_uuid, + forecast_version="0.0.0", + timestamp_utc=now, + ) + client.session.add(forecast) + client.session.commit() + + # Add forecast values for "Tomorrow" + # We need data that falls into "tomorrow" in IST. + # IST is UTC+5:30. + # Let's just generate 24 hours of data starting from now + 24h to be safe + start_future = now + dt.timedelta(hours=24) + + forecast_values = [] + for i in range(48): # 12 hours of data + start_utc = start_future + dt.timedelta(minutes=15 * i) + end_utc = start_utc + dt.timedelta(minutes=15) + + fv = ForecastValueSQL( + forecast_power_kw=123.45 + i, + forecast_uuid=forecast.forecast_uuid, + start_utc=start_utc, + end_utc=end_utc, + horizon_minutes=(start_utc - now).total_seconds() / 60, + ) + fv.ml_model = ml_model + forecast_values.append(fv) + + client.session.add_all(forecast_values) + client.session.commit() + + response = test_client.get( + "/wind/testID/forecast/csv", + params={"forecast_horizon": "day_ahead"}, + ) + + # The response should be successful even if there's no data for tomorrow + assert response.status_code == 200 + assert response.headers["content-type"] == "text/csv; charset=utf-8" + + # Parse CSV - it may be empty if no forecast data for tomorrow + csv_content = response.text + if csv_content.strip(): # Only validate if there's content + df = pd.read_csv(io.StringIO(csv_content)) + assert len(df.columns) == 3 + assert list(df.columns) == ["Date [IST]", "Time", "PowerMW"] From c73a950a7cdc024e97d4d5123398614315c6cfd8 Mon Sep 17 00:00:00 2001 From: Khusham Bansal Date: Fri, 26 Dec 2025 04:54:09 +0530 Subject: [PATCH 2/2] fix lint errors --- .../backends/quartzdb/test_csv_download.py | 39 +++++++++---------- 1 file changed, 19 insertions(+), 20 deletions(-) diff --git a/src/quartz_api/internal/backends/quartzdb/test_csv_download.py b/src/quartz_api/internal/backends/quartzdb/test_csv_download.py index 2bc6f94c..3790ffc3 100644 --- a/src/quartz_api/internal/backends/quartzdb/test_csv_download.py +++ b/src/quartz_api/internal/backends/quartzdb/test_csv_download.py @@ -1,11 +1,14 @@ """Test CSV download endpoint.""" +import datetime as dt import io import pathlib import pandas as pd import pytest from fastapi.testclient import TestClient +from pvsite_datamodel.read.model import get_or_create_model +from pvsite_datamodel.sqlmodels import ForecastSQL, ForecastValueSQL, LocationSQL from pyhocon import ConfigFactory from sqlalchemy.engine import Engine from sqlalchemy.orm import Session @@ -14,9 +17,6 @@ from quartz_api.internal import models from .client import Client -import datetime as dt -from pvsite_datamodel.read.model import get_or_create_model -from pvsite_datamodel.sqlmodels import ForecastSQL, ForecastValueSQL, LocationSQL @pytest.fixture() @@ -34,28 +34,27 @@ def test_client(client: Client) -> TestClient: # Load config and override routers to include 'regions' conf_path = pathlib.Path(__file__).parent.parent.parent.parent / "cmd" / "server.conf" conf = ConfigFactory.parse_file(conf_path.as_posix()) - + # Override config to include regions router override_conf = ConfigFactory.from_dict({"api": {"routers": "regions"}}) conf = override_conf.with_fallback(conf) - + # Create the FastAPI app app = _create_server(conf=conf) - + # Override the database dependency with our test client app.dependency_overrides[models.get_db_client] = lambda: client - + return TestClient(app) class TestCsvDownload: """Test CSV download functionality.""" + @pytest.mark.usefixtures("forecast_values_wind") def test_csv_download_success( self, test_client: TestClient, - client: Client, - forecast_values_wind: None, ) -> None: """Test successful CSV download from the API endpoint.""" # Make request to CSV download endpoint @@ -81,24 +80,24 @@ def test_csv_download_success( assert "PowerMW" in df.columns assert len(df) > 0 + @pytest.mark.usefixtures("forecast_values_wind") def test_csv_download_day_ahead( self, test_client: TestClient, client: Client, - forecast_values_wind: None, ) -> None: """Test CSV download with day_ahead forecast horizon.""" - + # 1. Setup: Create 'Day Ahead' forecast data (for tomorrow) # Get the wind site site = client.session.query(LocationSQL).filter(LocationSQL.asset_type == "wind").first() - + # Create a forecast timestamped now now = dt.datetime.now(tz=dt.UTC) - + # Create a model entry ml_model = get_or_create_model(client.session, "windnet_india_adjust") - + forecast = ForecastSQL( location_uuid=site.location_uuid, forecast_version="0.0.0", @@ -106,18 +105,18 @@ def test_csv_download_day_ahead( ) client.session.add(forecast) client.session.commit() - + # Add forecast values for "Tomorrow" - # We need data that falls into "tomorrow" in IST. + # We need data that falls into "tomorrow" in IST. # IST is UTC+5:30. # Let's just generate 24 hours of data starting from now + 24h to be safe start_future = now + dt.timedelta(hours=24) - + forecast_values = [] - for i in range(48): # 12 hours of data + for i in range(48): # 12 hours of data start_utc = start_future + dt.timedelta(minutes=15 * i) end_utc = start_utc + dt.timedelta(minutes=15) - + fv = ForecastValueSQL( forecast_power_kw=123.45 + i, forecast_uuid=forecast.forecast_uuid, @@ -127,7 +126,7 @@ def test_csv_download_day_ahead( ) fv.ml_model = ml_model forecast_values.append(fv) - + client.session.add_all(forecast_values) client.session.commit()