194 lines
6.4 KiB
Python
194 lines
6.4 KiB
Python
import asyncio
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
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
|