352 lines
14 KiB
Python
352 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Mapping
|
|
from zoneinfo import ZoneInfo
|
|
|
|
from arr_web.app import PortalApplication, RuntimeHealth, SessionLedger
|
|
from arr_web.company_jobs import (
|
|
COMPANY_NAMES,
|
|
PersistentCompanyReportCoordinator,
|
|
period_release_at,
|
|
period_to_as_of,
|
|
)
|
|
from arr_web.contracts import PortalError
|
|
from arr_web.downloads import ArtifactDescriptor, ControlledProjectArtifactReader
|
|
from company_reports.contracts import ENGLISH_MONTH_NAMES, RESULT_SCHEMA_VERSION
|
|
|
|
|
|
BANGKOK = ZoneInfo("Asia/Bangkok")
|
|
|
|
|
|
def decoded(response: Any) -> Dict[str, Any]:
|
|
return json.loads(response.body.decode("utf-8"))
|
|
|
|
|
|
class FixtureExecutor:
|
|
def __init__(self, project_root: Path, *, failed_company: str = "") -> None:
|
|
self.project_root = project_root
|
|
self.output_root = project_root / "outputs" / "company_reports"
|
|
self.failed_company = failed_company
|
|
self.calls: list[tuple[int, int, object]] = []
|
|
|
|
def run(self, year: int, month: int, as_of_date: object) -> Mapping[str, Any]:
|
|
self.calls.append((year, month, as_of_date))
|
|
final_day = 29 if year == 2028 and month == 2 else 31
|
|
if month in {4, 6, 9, 11}:
|
|
final_day = 30
|
|
elif month == 2 and year != 2028:
|
|
final_day = 28
|
|
companies = []
|
|
for index, company in enumerate(COMPANY_NAMES, 1):
|
|
counts = {"01-10": index, "11-20": index + 1, f"21-{final_day:02d}": 0}
|
|
if company == self.failed_company:
|
|
companies.append(
|
|
{
|
|
"company": company,
|
|
"status": "failed",
|
|
"row_count": 0,
|
|
"period_row_counts": counts,
|
|
"warnings": [],
|
|
"errors": [
|
|
{
|
|
"code": "COMPANY_REPORT_GROUP_CODE_NOT_FOUND",
|
|
"stage": "source",
|
|
"period": "01-10",
|
|
"record_ids": [100 + index],
|
|
}
|
|
],
|
|
}
|
|
)
|
|
continue
|
|
filename = f"{company}-{ENGLISH_MONTH_NAMES[month]}-{year}.xlsx"
|
|
path = self.output_root / f"{year:04d}" / f"{month:02d}" / "test" / filename
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
payload = f"safe workbook fixture {company}".encode("utf-8")
|
|
path.write_bytes(payload)
|
|
companies.append(
|
|
{
|
|
"company": company,
|
|
"status": "success",
|
|
"row_count": sum(counts.values()),
|
|
"period_row_counts": counts,
|
|
"version_no": index,
|
|
"warnings": (
|
|
[
|
|
{
|
|
"code": "COMPANY_REPORT_MULTI_PRICE_REVIEW",
|
|
"stage": "pricing",
|
|
"period": "01-10",
|
|
"record_ids": [index],
|
|
}
|
|
]
|
|
if company == "QBD"
|
|
else []
|
|
),
|
|
"errors": [],
|
|
"artifact": {
|
|
"filename": filename,
|
|
"storage_key": path.relative_to(self.project_root).as_posix(),
|
|
"sha256": hashlib.sha256(payload).hexdigest(),
|
|
"semantic_sha256": "a" * 64,
|
|
},
|
|
}
|
|
)
|
|
successful = len(COMPANY_NAMES) - (1 if self.failed_company else 0)
|
|
status = "success" if successful == len(COMPANY_NAMES) else "partial_failure"
|
|
return {
|
|
"schema_version": RESULT_SCHEMA_VERSION,
|
|
"status": status,
|
|
"report_year": year,
|
|
"report_month": month,
|
|
"as_of_date": str(as_of_date),
|
|
"requested_companies": list(COMPANY_NAMES),
|
|
"companies": companies,
|
|
}
|
|
|
|
|
|
class BlockingFailureExecutor:
|
|
def __init__(self) -> None:
|
|
self.started = threading.Event()
|
|
self.release = threading.Event()
|
|
|
|
def run(self, year: int, month: int, as_of_date: object) -> Mapping[str, Any]:
|
|
self.started.set()
|
|
self.release.wait(timeout=3)
|
|
counts = {"01-10": 0, "11-20": 0, "21-31": 0}
|
|
return {
|
|
"schema_version": RESULT_SCHEMA_VERSION,
|
|
"status": "failed",
|
|
"report_year": year,
|
|
"report_month": month,
|
|
"as_of_date": str(as_of_date),
|
|
"requested_companies": list(COMPANY_NAMES),
|
|
"companies": [
|
|
{
|
|
"company": company,
|
|
"status": "failed",
|
|
"row_count": 0,
|
|
"period_row_counts": counts,
|
|
"warnings": [],
|
|
"errors": [
|
|
{
|
|
"code": "COMPANY_REPORT_SOURCE_VERSION_MISSING",
|
|
"stage": "source",
|
|
"period": "2026-07",
|
|
"record_ids": [],
|
|
}
|
|
],
|
|
}
|
|
for company in COMPANY_NAMES
|
|
],
|
|
}
|
|
|
|
|
|
def wait_terminal(
|
|
coordinator: PersistentCompanyReportCoordinator,
|
|
job_id: str,
|
|
timeout: float = 3.0,
|
|
) -> Dict[str, Any]:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
value = coordinator.get_job(job_id)
|
|
if value["state"] in {"succeeded", "partial_failure", "failed"}:
|
|
return value
|
|
time.sleep(0.01)
|
|
raise AssertionError("company report job did not finish")
|
|
|
|
|
|
class CompanyReportCoordinatorTests(unittest.TestCase):
|
|
def test_period_mapping_and_bangkok_release_instants(self) -> None:
|
|
self.assertEqual(str(period_to_as_of("2028-02", "21-month-end")), "2028-02-29")
|
|
self.assertEqual(
|
|
period_release_at("2026-07", "01-10").isoformat(),
|
|
"2026-07-11T00:00:00+07:00",
|
|
)
|
|
self.assertEqual(
|
|
period_release_at("2026-12", "21-month-end").isoformat(),
|
|
"2027-01-01T00:00:00+07:00",
|
|
)
|
|
|
|
def test_success_is_private_and_download_is_reverified(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
executor = FixtureExecutor(root)
|
|
coordinator = PersistentCompanyReportCoordinator(
|
|
root,
|
|
root / "outputs" / "company_reports",
|
|
executor,
|
|
now=lambda: datetime(2026, 8, 1, tzinfo=BANGKOK),
|
|
)
|
|
try:
|
|
created = coordinator.create("2026-07", "11-20")
|
|
job = wait_terminal(coordinator, created["job_id"])
|
|
self.assertEqual(job["state"], "succeeded")
|
|
self.assertEqual(job["requested_companies"], list(COMPANY_NAMES))
|
|
self.assertEqual(len(job["company_results"]), 5)
|
|
self.assertEqual(job["company_results"][1]["warnings"][0]["record_ids"], [2])
|
|
public_text = json.dumps(job, ensure_ascii=False)
|
|
self.assertNotIn("storage_key", public_text)
|
|
self.assertNotIn("sha256", public_text)
|
|
descriptor = coordinator.resolve_download(created["job_id"], "QBD")
|
|
self.assertEqual(descriptor.file_kind, "company_ten_day_xlsx")
|
|
body = ControlledProjectArtifactReader(root).read(descriptor)
|
|
self.assertIn(b"safe workbook fixture", body)
|
|
finally:
|
|
coordinator.close()
|
|
|
|
def test_partial_failure_keeps_successful_company_downloads(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
coordinator = PersistentCompanyReportCoordinator(
|
|
root,
|
|
root / "outputs" / "company_reports",
|
|
FixtureExecutor(root, failed_company="HanaTour"),
|
|
now=lambda: datetime(2026, 8, 1, tzinfo=BANGKOK),
|
|
)
|
|
try:
|
|
created = coordinator.create("2026-07", "01-10")
|
|
job = wait_terminal(coordinator, created["job_id"])
|
|
self.assertEqual(job["state"], "partial_failure")
|
|
self.assertEqual(set(job["downloads"]), set(COMPANY_NAMES) - {"HanaTour"})
|
|
with self.assertRaises(PortalError) as caught:
|
|
coordinator.resolve_download(created["job_id"], "HanaTour")
|
|
self.assertEqual(caught.exception.status, 404)
|
|
finally:
|
|
coordinator.close()
|
|
|
|
def test_unreleased_and_duplicate_active_jobs_are_rejected(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
blocked = PersistentCompanyReportCoordinator(
|
|
root,
|
|
root / "outputs" / "company_reports",
|
|
FixtureExecutor(root),
|
|
now=lambda: datetime(2026, 7, 10, 23, 59, tzinfo=BANGKOK),
|
|
)
|
|
try:
|
|
with self.assertRaises(PortalError) as caught:
|
|
blocked.create("2026-07", "01-10")
|
|
self.assertEqual(caught.exception.code, "COMPANY_REPORT_PERIOD_NOT_OPEN")
|
|
finally:
|
|
blocked.close()
|
|
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
executor = BlockingFailureExecutor()
|
|
coordinator = PersistentCompanyReportCoordinator(
|
|
root,
|
|
root / "outputs" / "company_reports",
|
|
executor,
|
|
now=lambda: datetime(2026, 8, 1, tzinfo=BANGKOK),
|
|
)
|
|
try:
|
|
first = coordinator.create("2026-07", "01-10")
|
|
self.assertTrue(executor.started.wait(timeout=1))
|
|
with self.assertRaises(PortalError) as caught:
|
|
coordinator.create("2026-07", "01-10")
|
|
self.assertEqual(caught.exception.code, "COMPANY_REPORT_ALREADY_RUNNING")
|
|
executor.release.set()
|
|
self.assertEqual(wait_terminal(coordinator, first["job_id"])["state"], "failed")
|
|
finally:
|
|
executor.release.set()
|
|
coordinator.close()
|
|
|
|
|
|
class FakeCompanyCoordinator:
|
|
def __init__(self) -> None:
|
|
self.created: list[tuple[str, str]] = []
|
|
self.descriptor = ArtifactDescriptor(
|
|
"company_ten_day_xlsx",
|
|
"QBD-July-2026.xlsx",
|
|
"outputs/company_reports/QBD-July-2026.xlsx",
|
|
"a" * 64,
|
|
4,
|
|
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
|
)
|
|
|
|
def create(self, report_month: str, period: str) -> Dict[str, Any]:
|
|
self.created.append((report_month, period))
|
|
return {"job_id": "a" * 32, "state": "queued"}
|
|
|
|
def list_jobs(self, report_month: object = None, limit: int = 100) -> list[Dict[str, Any]]:
|
|
return [{"job_id": "a" * 32, "report_month": report_month, "limit": limit}]
|
|
|
|
def get_job(self, job_id: str) -> Dict[str, Any]:
|
|
return {"job_id": job_id, "state": "succeeded"}
|
|
|
|
def resolve_download(self, job_id: str, company: str) -> ArtifactDescriptor:
|
|
return self.descriptor
|
|
|
|
|
|
class FakeReader:
|
|
def read(self, descriptor: ArtifactDescriptor) -> bytes:
|
|
return b"xlsx"
|
|
|
|
|
|
class CompanyReportRouteTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.coordinator = FakeCompanyCoordinator()
|
|
self.app = PortalApplication(
|
|
company_reports=self.coordinator,
|
|
artifact_reader=FakeReader(),
|
|
health=RuntimeHealth(False, False, False, True, True),
|
|
sessions=SessionLedger(),
|
|
)
|
|
|
|
def session_headers(self) -> Dict[str, str]:
|
|
response = self.app.handle("GET", "/api/session", {})
|
|
payload = decoded(response)
|
|
cookie = response.headers["Set-Cookie"].split(";", 1)[0]
|
|
return {"Cookie": cookie, "X-ARR-CSRF": payload["data"]["csrf_token"]}
|
|
|
|
def test_create_is_csrf_protected_and_strict(self) -> None:
|
|
body = b'{"report_month":"2026-07","period":"01-10"}'
|
|
self.assertEqual(
|
|
self.app.handle("POST", "/api/company-reports/jobs", {}, body).status,
|
|
403,
|
|
)
|
|
response = self.app.handle(
|
|
"POST",
|
|
"/api/company-reports/jobs",
|
|
{**self.session_headers(), "Content-Type": "application/json"},
|
|
body,
|
|
)
|
|
self.assertEqual(response.status, 202)
|
|
self.assertEqual(self.coordinator.created, [("2026-07", "01-10")])
|
|
invalid = self.app.handle(
|
|
"POST",
|
|
"/api/company-reports/jobs",
|
|
{**self.session_headers(), "Content-Type": "application/json"},
|
|
b'{"report_month":"2026-07","period":"01-10","companies":[]}',
|
|
)
|
|
self.assertEqual(invalid.status, 400)
|
|
self.assertEqual(decoded(invalid)["error"]["code"], "COMPANY_REPORT_REQUEST_INVALID")
|
|
|
|
def test_list_detail_and_download_routes(self) -> None:
|
|
listing = self.app.handle(
|
|
"GET", "/api/company-reports/jobs?month=2026-07&limit=20", {}
|
|
)
|
|
self.assertEqual(decoded(listing)["data"][0]["report_month"], "2026-07")
|
|
detail = self.app.handle("GET", f"/api/company-reports/jobs/{'a' * 32}", {})
|
|
self.assertEqual(decoded(detail)["data"]["state"], "succeeded")
|
|
download = self.app.handle(
|
|
"GET",
|
|
f"/api/company-reports/jobs/{'a' * 32}/downloads/QBD",
|
|
{},
|
|
)
|
|
self.assertEqual(download.status, 200)
|
|
self.assertEqual(download.body, b"xlsx")
|
|
self.assertNotIn("outputs/company_reports", str(download.headers))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|