solorpower/api_server/tests/test_upload.py

import io
import unittest
from unittest.mock import patch

import pandas as pd
from fastapi.testclient import TestClient

from app.core.database import get_db
from app.main import app
from app.routers import upload


class FakeResponse:
    def __init__(self, data=None):
        self.data = data or []


class FakeRpc:
    def __init__(self, db, params):
        self.db = db
        self.params = params

    def execute(self):
        self.db.saved["daily_stats"] = self.params["p_records"]
        self.db.rpc_params = self.params
        return FakeResponse(self.params["p_records"])


class FakeQuery:
    def __init__(self, db, table_name):
        self.db = db
        self.table_name = table_name
        self.action = None
        self.payload = None
        self.filters = {}

    def select(self, _columns):
        self.action = "select"
        return self

    def eq(self, column, value):
        self.filters[column] = value
        return self

    def limit(self, _count):
        return self

    def upsert(self, payload, on_conflict=None):
        self.action = "upsert"
        self.payload = payload
        self.on_conflict = on_conflict
        return self

    def execute(self):
        if self.action == "select" and self.table_name == "plants":
            if self.filters.get("id") != "plant-a":
                return FakeResponse()
            return FakeResponse([{
                "id": "plant-a",
                "name": "테스트 발전소",
                "capacity": 100.0,
            }])

        if self.action == "upsert":
            self.db.saved[self.table_name] = self.payload
            return FakeResponse(self.payload)

        raise AssertionError(f"unexpected query: {self.table_name} {self.action}")


class FakeDb:
    def __init__(self):
        self.saved = {}
        self.rpc_params = None

    def table(self, table_name):
        return FakeQuery(self, table_name)

    def rpc(self, function_name, params):
        if function_name != "upsert_daily_stats":
            raise AssertionError(function_name)
        return FakeRpc(self, params)


def excel_bytes(rows):
    output = io.BytesIO()
    pd.DataFrame(rows).to_excel(output, index=False)
    return output.getvalue()


class UploadApiTest(unittest.TestCase):
    def setUp(self):
        self.db = FakeDb()
        app.dependency_overrides[get_db] = lambda: self.db
        self.client = TestClient(app)

    def tearDown(self):
        app.dependency_overrides.clear()

    def post_daily(self, content, filename="daily.xlsx", plant_id="plant-a"):
        return self.client.post(
            f"/plants/{plant_id}/upload",
            files={"file": (filename, content, "application/octet-stream")},
        )

    def post_monthly(self, content, filename="monthly.xlsx", plant_id="plant-a"):
        return self.client.post(
            f"/plants/{plant_id}/upload/monthly",
            files={"file": (filename, content, "application/octet-stream")},
        )

    def test_daily_sample_is_saved(self):
        response = self.post_daily(excel_bytes([
            {"date": "2026-08-05", "generation": 123.4},
        ]))

        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json()["saved_count"], 1)
        self.assertEqual(self.db.saved["daily_stats"][0]["total_generation"], 123.4)
        self.assertEqual(self.db.rpc_params["p_source"], "excel_daily")
        self.assertTrue(self.db.rpc_params["p_allow_decrease"])

    def test_monthly_sample_with_merged_year_is_saved(self):
        response = self.post_monthly(excel_bytes([
            {"year": 2026, "month": "1월", "kwh": "1,234.5"},
            {"year": None, "month": "2월", "kwh": 987.6},
        ]))

        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json()["saved_count"], 2)
        self.assertEqual(self.db.saved["monthly_stats"][1]["month"], "2026-02")
        self.assertEqual(self.db.saved["monthly_stats"][1]["source"], "excel_monthly")
        self.assertEqual(self.db.saved["monthly_stats"][1]["last_date"], "2026-02-28")

    def test_wrong_extension_returns_400(self):
        response = self.post_daily(b"text", filename="daily.csv")
        self.assertEqual(response.status_code, 400)

    def test_empty_file_returns_400(self):
        response = self.post_daily(b"")
        self.assertEqual(response.status_code, 400)

    def test_missing_columns_returns_400(self):
        response = self.post_daily(excel_bytes([{"wrong": 1}]))
        self.assertEqual(response.status_code, 400)

    def test_invalid_date_returns_400(self):
        response = self.post_daily(excel_bytes([
            {"date": "not-a-date", "generation": 1},
        ]))
        self.assertEqual(response.status_code, 400)

    def test_negative_generation_returns_400(self):
        response = self.post_daily(excel_bytes([
            {"date": "2026-08-05", "generation": -1},
        ]))
        self.assertEqual(response.status_code, 400)

    def test_unknown_plant_returns_404(self):
        response = self.post_daily(
            excel_bytes([{"date": "2026-08-05", "generation": 1}]),
            plant_id="missing",
        )
        self.assertEqual(response.status_code, 404)

    def test_file_size_limit_returns_413(self):
        response = self.post_daily(b"x" * (upload.MAX_UPLOAD_BYTES + 1))
        self.assertEqual(response.status_code, 413)

    def test_row_limit_returns_400(self):
        with patch.object(upload, "MAX_UPLOAD_ROWS", 1):
            response = self.post_daily(excel_bytes([
                {"date": "2026-08-05", "generation": 1},
                {"date": "2026-08-06", "generation": 2},
            ]))
        self.assertEqual(response.status_code, 400)


if __name__ == "__main__":
    unittest.main()