TJWaterServerBinary/tests/api/test_model_import_endpoints.py

import asyncio
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
from uuid import uuid4

import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient

from app.api.v1.endpoints import model_import
from app.auth.metadata_dependencies import (
    get_current_metadata_admin,
    get_metadata_repository,
)
from app.infra.db.metadb.repositories.metadata_repository import ProjectDbRouting
from app.infra.db.project_routing import get_project_pgconn_string
from app.native.wndb.core.database import MaterializedViewRefreshAfterCommitError
from tests.conftest import build_test_app


VALID_INP = b"[TITLE]\nDesktop model\n[JUNCTIONS]\n;ID Elev Demand\n"


def _client(*, admin=None, repo=None) -> TestClient:
    app = build_test_app(model_import.router, "/api/v1")
    if admin is not None:
        app.dependency_overrides[get_current_metadata_admin] = lambda: admin
    if repo is not None:
        app.dependency_overrides[get_metadata_repository] = lambda: repo
    return TestClient(app)


def test_system_admin_can_import_model_without_project_membership(
    monkeypatch,
):
    project_id = uuid4()
    project = SimpleNamespace(id=project_id, code="demo", status="active")
    repo = SimpleNamespace(
        session=object(),
        get_project_by_id=AsyncMock(return_value=project),
    )
    admin = SimpleNamespace(id=uuid4(), role="admin", is_superuser=False)
    monkeypatch.setattr(
        model_import,
        "_run_uploaded_inp",
        AsyncMock(return_value="imported"),
    )
    monkeypatch.setattr(model_import, "log_audit_event", AsyncMock())
    client = _client(admin=admin, repo=repo)

    response = client.post(
        f"/api/v1/admin/projects/{project_id}/model-imports",
        files={"file": ("desktop-model.inp", VALID_INP)},
    )

    assert response.status_code == 200
    assert response.json()["project_id"] == str(project_id)
    assert response.json()["result"] == "imported"
    repo.get_project_by_id.assert_awaited_once_with(project_id)
    model_import.log_audit_event.assert_awaited_once()


def test_non_admin_is_denied_model_import():
    def deny_admin():
        raise HTTPException(status_code=403, detail="Admin access required")

    app = build_test_app(model_import.router, "/api/v1")
    app.dependency_overrides[get_current_metadata_admin] = deny_admin
    client = TestClient(app)

    response = client.post(
        f"/api/v1/admin/projects/{uuid4()}/model-imports",
        files={"file": ("desktop-model.inp", VALID_INP)},
    )

    assert response.status_code == 403
    assert response.json()["detail"] == "Admin access required"


def test_model_import_rejects_non_inp_file(monkeypatch):
    project_id = uuid4()
    repo = SimpleNamespace(
        session=object(),
        get_project_by_id=AsyncMock(
            return_value=SimpleNamespace(
                id=project_id,
                code="demo",
                status="active",
            )
        ),
    )
    monkeypatch.setattr(model_import, "log_audit_event", AsyncMock())
    client = _client(
        admin=SimpleNamespace(id=uuid4(), role="admin", is_superuser=False),
        repo=repo,
    )

    response = client.post(
        f"/api/v1/admin/projects/{project_id}/model-imports",
        files={"file": ("desktop-model.txt", VALID_INP)},
    )

    assert response.status_code == 400
    assert response.json()["detail"] == "Only .inp model files are accepted"
    model_import.log_audit_event.assert_not_awaited()


def test_model_update_uses_project_business_routing(monkeypatch):
    project_id = uuid4()
    project = SimpleNamespace(id=project_id, code="demo", status="active")
    repo = SimpleNamespace(
        session=object(),
        get_project_by_id=AsyncMock(return_value=project),
        get_project_db_routing=AsyncMock(
            return_value=ProjectDbRouting(
                project_id=project_id,
                db_role="biz_data",
                db_type="postgresql",
                dsn="postgresql://user:password@biz.example/routed_business",
                pool_min_size=1,
                pool_max_size=5,
            )
        ),
    )
    captured: dict[str, str] = {}

    async def fake_apply_model_update(content: bytes, project_code: str) -> None:
        assert content == VALID_INP
        captured["project_code"] = project_code
        captured["dsn"] = get_project_pgconn_string(project_code)

    monkeypatch.setattr(model_import, "_apply_model_update", fake_apply_model_update)
    monkeypatch.setattr(model_import, "log_audit_event", AsyncMock())
    client = _client(
        admin=SimpleNamespace(id=uuid4(), role="admin", is_superuser=False),
        repo=repo,
    )

    response = client.patch(
        f"/api/v1/admin/projects/{project_id}/model-imports",
        files={"file": ("desktop-model.inp", VALID_INP)},
    )

    assert response.status_code == 200
    assert captured == {
        "project_code": "demo",
        "dsn": "postgresql://user:password@biz.example/routed_business",
    }
    repo.get_project_db_routing.assert_awaited_once_with(project_id, "biz_data")


def test_gb18030_upload_is_normalized_to_utf8() -> None:
    content = "[TITLE]\n天津供水\n[JUNCTIONS]\n".encode("gb18030")

    class FakeUpload:
        filename = "model.inp"

        async def read(self, _limit: int) -> bytes:
            return content

    normalized, filename = asyncio.run(model_import._read_upload(FakeUpload()))

    assert filename == "model.inp"
    assert normalized.decode("utf-8") == "[TITLE]\n天津供水\n[JUNCTIONS]\n"


def test_model_update_runs_blocking_import_in_threadpool(monkeypatch) -> None:
    calls: list[tuple[object, tuple[object, ...]]] = []

    async def fake_threadpool(function, *args):
        calls.append((function, args))

    monkeypatch.setattr(model_import, "run_in_threadpool", fake_threadpool)

    asyncio.run(model_import._update_from_inp(b"[TITLE]\n", "demo"))

    assert calls == [(model_import._update_from_inp_sync, (b"[TITLE]\n", "demo"))]


def test_committed_refresh_failure_is_not_wrapped_as_retryable_500(
    monkeypatch,
) -> None:
    error = MaterializedViewRefreshAfterCommitError("demo")

    async def fail_update(_content: bytes, _project_code: str) -> None:
        raise error

    monkeypatch.setattr(model_import, "_update_from_inp", fail_update)

    with pytest.raises(MaterializedViewRefreshAfterCommitError) as exc_info:
        asyncio.run(model_import._apply_model_update(b"[TITLE]\n", "demo"))

    assert exc_info.value is error


def test_project_provision_creates_metadata_only_after_infrastructure(monkeypatch):
    project_id = uuid4()
    now = datetime(2026, 1, 1, tzinfo=timezone.utc)
    project = SimpleNamespace(
        id=project_id,
        name="Demo",
        code="demo",
        description=None,
        gs_workspace="demo",
        map_extent={"bbox": [1.0, 2.0, 3.0, 4.0], "zoom": 15},
        status="active",
        created_at=now,
        updated_at=now,
    )
    infrastructure = SimpleNamespace(
        map_bbox=(1.0, 2.0, 3.0, 4.0),
        business_dsn="postgresql://business/demo",
        timescale_dsn="postgresql://timescale/demo",
        model_template="demo_template",
        layers=("junctions", "pipes"),
    )
    repo = SimpleNamespace(
        session=SimpleNamespace(rollback=AsyncMock()),
        get_project_by_code=AsyncMock(return_value=None),
        create_provisioned_project=AsyncMock(return_value=project),
    )
    monkeypatch.setattr(
        model_import,
        "_run_uploaded_inp",
        AsyncMock(return_value='{"simulation_result":"successful"}'),
    )
    monkeypatch.setattr(
        model_import,
        "is_database_encryption_configured",
        lambda: True,
    )

    async def fake_threadpool(function, *args, **kwargs):
        assert function is model_import._provision_from_inp_sync
        return infrastructure

    monkeypatch.setattr(model_import, "run_in_threadpool", fake_threadpool)
    monkeypatch.setattr(model_import, "log_audit_event", AsyncMock())
    client = _client(
        admin=SimpleNamespace(id=uuid4(), role="admin", is_superuser=True),
        repo=repo,
    )

    response = client.post(
        "/api/v1/admin/project-provisions",
        data={"name": "Demo", "code": "demo", "map_zoom": "15"},
        files={"file": ("model.inp", VALID_INP)},
    )

    assert response.status_code == 201
    assert response.json()["model_template_database"] == "demo_template"
    repo.create_provisioned_project.assert_awaited_once()
    assert repo.create_provisioned_project.await_args.kwargs["business_dsn"] == infrastructure.business_dsn
    model_import.log_audit_event.assert_awaited_once()


def test_project_provision_cleans_infrastructure_when_metadata_commit_fails(monkeypatch):
    cleanup = Mock(return_value=[])
    infrastructure = SimpleNamespace(
        map_bbox=(1.0, 2.0, 3.0, 4.0),
        business_dsn="postgresql://business/demo",
        timescale_dsn="postgresql://timescale/demo",
        model_template="demo_template",
        layers=("junctions", "pipes"),
        cleanup=cleanup,
    )
    repo = SimpleNamespace(
        session=SimpleNamespace(rollback=AsyncMock()),
        get_project_by_code=AsyncMock(return_value=None),
        create_provisioned_project=AsyncMock(side_effect=RuntimeError("metadata down")),
    )
    monkeypatch.setattr(
        model_import,
        "_run_uploaded_inp",
        AsyncMock(return_value='{"simulation_result":"successful"}'),
    )
    monkeypatch.setattr(
        model_import,
        "is_database_encryption_configured",
        lambda: True,
    )

    async def fake_threadpool(function, *args, **kwargs):
        if function is model_import._provision_from_inp_sync:
            return infrastructure
        return function(*args, **kwargs)

    monkeypatch.setattr(model_import, "run_in_threadpool", fake_threadpool)
    monkeypatch.setattr(model_import, "log_audit_event", AsyncMock())
    client = _client(
        admin=SimpleNamespace(id=uuid4(), role="admin", is_superuser=True),
        repo=repo,
    )

    response = client.post(
        "/api/v1/admin/project-provisions",
        data={"name": "Demo", "code": "demo"},
        files={"file": ("model.inp", VALID_INP)},
    )

    assert response.status_code == 503
    repo.session.rollback.assert_awaited_once()
    cleanup.assert_called_once_with()
    model_import.log_audit_event.assert_not_awaited()