diff --git a/app/algorithms/simulation/scenarios.py b/app/algorithms/simulation/scenarios.py index 9890fc8..28f9954 100644 --- a/app/algorithms/simulation/scenarios.py +++ b/app/algorithms/simulation/scenarios.py @@ -72,6 +72,7 @@ def burst_analysis( modify_variable_pump_pattern: dict[str, list] = None, modify_valve_opening: dict[str, float] = None, scheme_name: str = None, + username: str | None = None, ) -> None: """ 爆管模拟 @@ -86,6 +87,9 @@ def burst_analysis( :param scheme_name: 方案名称 :return: """ + if not username: + raise ValueError("username is required when storing burst analysis scheme") + scheme_detail: dict = { "burst_ID": burst_ID, "burst_size": burst_size, @@ -211,7 +215,7 @@ def burst_analysis( name=name, scheme_name=scheme_name, scheme_type="burst_analysis", - username="admin", + username=username, scheme_start_time=modify_pattern_start_time, scheme_detail=scheme_detail, ) @@ -311,6 +315,7 @@ def flushing_analysis( drainage_node_ID: str = None, flushing_flow: float = 0, scheme_name: str = None, + username: str | None = None, ) -> None: """ 管道冲洗模拟 @@ -323,6 +328,9 @@ def flushing_analysis( :param scheme_name: 方案名称 :return: """ + if not username: + raise ValueError("username is required when storing flushing analysis scheme") + scheme_detail: dict = { "duration": modify_total_duration, "valve_opening": modify_valve_opening, @@ -455,7 +463,7 @@ def flushing_analysis( name=name, scheme_name=scheme_name, scheme_type="flushing_analysis", - username="admin", + username=username, scheme_start_time=modify_pattern_start_time, scheme_detail=scheme_detail, ) @@ -473,6 +481,7 @@ def contaminant_simulation( concentration: float, # 污染源浓度,单位mg/L scheme_name: str = None, source_pattern: str = None, # 污染源时间变化模式名称 + username: str | None = None, ) -> None: """ 污染模拟 @@ -486,6 +495,9 @@ def contaminant_simulation( :param scheme_name: 方案名称 :return: """ + if not username: + raise ValueError("username is required when storing contaminant analysis scheme") + scheme_detail: dict = { "source": source, "concentration": concentration, @@ -608,7 +620,7 @@ def contaminant_simulation( name=name, scheme_name=scheme_name, scheme_type="contaminant_analysis", - username="admin", + username=username, scheme_start_time=modify_pattern_start_time, scheme_detail=scheme_detail, ) diff --git a/app/api/v1/endpoints/simulation.py b/app/api/v1/endpoints/simulation.py index 3a758cd..2295aa8 100644 --- a/app/api/v1/endpoints/simulation.py +++ b/app/api/v1/endpoints/simulation.py @@ -4,8 +4,9 @@ import json import os import shutil import threading -from fastapi import APIRouter, HTTPException, File, UploadFile, Query, Path, Body +from fastapi import APIRouter, Depends, HTTPException, File, UploadFile, Query, Path, Body from fastapi.responses import PlainTextResponse +from app.auth.keycloak_dependencies import get_current_keycloak_username import app.services.simulation as simulation import app.services.globals as globals from app.services.tjnetwork import ( @@ -209,6 +210,7 @@ async def fastapi_burst_analysis( burst_size: list[float] = Query(..., description="对应各爆管点的爆管流量大小列表(L/s)"), modify_total_duration: int = Query(..., description="模拟总时长(秒)"), scheme_name: str = Query(..., description="分析方案名称"), + username: str = Depends(get_current_keycloak_username), ) -> str: """ 爆管分析(高级版本) @@ -229,6 +231,7 @@ async def fastapi_burst_analysis( burst_size=burst_size, modify_total_duration=modify_total_duration, scheme_name=scheme_name, + username=username, ) return "success" @@ -314,6 +317,7 @@ async def fastapi_flushing_analysis( flush_flow: float = Query(0, description="冲洗流量(L/s),0表示自动计算"), duration: int | None = Query(None, description="模拟持续时间(秒),默认900秒"), scheme_name: str = Query(..., description="冲洗方案名称"), + username: str = Depends(get_current_keycloak_username), ) -> str: """ 冲洗分析(高级版本) @@ -340,6 +344,7 @@ async def fastapi_flushing_analysis( drainage_node_ID=drainage_node_ID, flushing_flow=flush_flow, scheme_name=scheme_name, + username=username, ) return result or "success" @@ -354,6 +359,7 @@ async def fastapi_contaminant_simulation( duration: int = Query(..., description="模拟持续时间(秒)"), scheme_name: str = Query(..., description="模拟方案名称"), pattern: str | None = Query(None, description="污染源模式ID(可选)"), + username: str = Depends(get_current_keycloak_username), ) -> str: """ 污染物模拟 @@ -376,6 +382,7 @@ async def fastapi_contaminant_simulation( source=source, concentration=concentration, source_pattern=pattern, + username=username, ) return result or "success" diff --git a/tests/api/test_simulation_endpoints.py b/tests/api/test_simulation_endpoints.py index e022f92..c5c53b3 100644 --- a/tests/api/test_simulation_endpoints.py +++ b/tests/api/test_simulation_endpoints.py @@ -108,6 +108,12 @@ def _load_simulation_module(monkeypatch): ) +def _build_authenticated_client(module) -> TestClient: + app = build_test_app(module.router, "/api/v1") + app.dependency_overrides[module.get_current_keycloak_username] = lambda: "alice" + return TestClient(app) + + def test_run_project_endpoint_returns_plain_text(monkeypatch): module = _load_simulation_module(monkeypatch) monkeypatch.setattr(module, "run_project", lambda network: f"report::{network}") @@ -339,6 +345,33 @@ def test_valve_close_endpoint_passes_scheme_name(monkeypatch): } +def test_burst_endpoint_passes_current_username(monkeypatch): + module = _load_simulation_module(monkeypatch) + captured = {} + + def fake_burst_analysis(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(module, "burst_analysis", fake_burst_analysis) + client = _build_authenticated_client(module) + + response = client.get( + "/api/v1/burst_analysis/", + params={ + "network": "demo", + "modify_pattern_start_time": "2025-01-02T03:04:05+08:00", + "burst_ID": ["P1"], + "burst_size": [10.0], + "modify_total_duration": 900, + "scheme_name": "burst_case_01", + }, + ) + + assert response.status_code == 200 + assert response.text == '"success"' + assert captured["username"] == "alice" + + def test_flushing_endpoint_passes_required_scheme_name(monkeypatch): module = _load_simulation_module(monkeypatch) captured = {} @@ -348,7 +381,7 @@ def test_flushing_endpoint_passes_required_scheme_name(monkeypatch): return "ok" monkeypatch.setattr(module, "flushing_analysis", fake_flushing_analysis) - client = TestClient(build_test_app(module.router, "/api/v1")) + client = _build_authenticated_client(module) response = client.get( "/api/v1/flushing_analysis/", @@ -374,12 +407,41 @@ def test_flushing_endpoint_passes_required_scheme_name(monkeypatch): "drainage_node_ID": "N1", "flushing_flow": 100.0, "scheme_name": "flush_case_01", + "username": "alice", } +def test_contaminant_endpoint_passes_current_username(monkeypatch): + module = _load_simulation_module(monkeypatch) + captured = {} + + def fake_contaminant_simulation(**kwargs): + captured.update(kwargs) + return "ok" + + monkeypatch.setattr(module, "contaminant_simulation", fake_contaminant_simulation) + client = _build_authenticated_client(module) + + response = client.get( + "/api/v1/contaminant_simulation/", + params={ + "network": "demo", + "start_time": "2025-01-02T03:04:05+08:00", + "source": "N1", + "concentration": 10.0, + "duration": 900, + "scheme_name": "contaminant_case_01", + }, + ) + + assert response.status_code == 200 + assert response.text == "ok" + assert captured["username"] == "alice" + + def test_contaminant_endpoint_requires_scheme_name(monkeypatch): module = _load_simulation_module(monkeypatch) - client = TestClient(build_test_app(module.router, "/api/v1")) + client = _build_authenticated_client(module) response = client.get( "/api/v1/contaminant_simulation/",