81 lines
3.1 KiB
Python
81 lines
3.1 KiB
Python
from unittest.mock import Mock
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.exceptions import RequestValidationError
|
|
from fastapi.testclient import TestClient
|
|
|
|
from modelforge_api.api.authorization import (
|
|
Principal,
|
|
PrincipalRole,
|
|
authenticate_operator_principal,
|
|
)
|
|
from modelforge_api.api.routes.lab import get_lab_service, router
|
|
from modelforge_api.domain.lab import LabInputSummary, LabLimits, LabPreflightResponse
|
|
from modelforge_api.main import validation_error
|
|
|
|
|
|
def request_body() -> dict:
|
|
return {
|
|
"target": {name: str(uuid4()) for name in (
|
|
"deployment_candidate_id", "artifact_set_id", "runtime_profile_id",
|
|
"compute_node_id", "execution_approval_id",
|
|
)},
|
|
"input": {"capability": "rag.embedding", "payload": {"input": "private-input-canary"}},
|
|
}
|
|
|
|
|
|
def harness(role: PrincipalRole | None) -> tuple[TestClient, Mock]:
|
|
app = FastAPI()
|
|
app.include_router(router)
|
|
app.add_exception_handler(RequestValidationError, validation_error) # type: ignore[arg-type]
|
|
|
|
def principal() -> Principal:
|
|
if role is None:
|
|
raise HTTPException(status_code=401, detail="Authentication required")
|
|
return Principal("test-operator", role, "test")
|
|
|
|
service = Mock()
|
|
service.preflight.return_value = LabPreflightResponse(
|
|
input_summary=LabInputSummary(sha256="a" * 64, item_count=1, encoded_bytes=40),
|
|
limits=LabLimits(), blockers=["CANDIDATE_NOT_FOUND"], checks_passed=False,
|
|
)
|
|
app.dependency_overrides[authenticate_operator_principal] = principal
|
|
app.dependency_overrides[get_lab_service] = lambda: service
|
|
return TestClient(app), service
|
|
|
|
|
|
@pytest.mark.parametrize("role,status", [
|
|
(None, 401), (PrincipalRole.VIEWER, 403),
|
|
(PrincipalRole.OPERATOR, 200), (PrincipalRole.ADMIN, 200),
|
|
])
|
|
def test_preflight_requires_operator_without_execution(role: PrincipalRole | None, status: int) -> None:
|
|
client, service = harness(role)
|
|
response = client.post("/api/v1/lab/preflight", json=request_body())
|
|
assert response.status_code == status
|
|
assert service.preflight.call_count == int(status == 200)
|
|
if status == 200:
|
|
assert response.json()["execution_started"] is False
|
|
assert response.json()["permission_granted"] is False
|
|
assert "private-input-canary" not in response.text
|
|
|
|
|
|
@pytest.mark.parametrize("location", ["discriminator", "payload", "extra"])
|
|
def test_validation_never_reflects_payload_or_discriminator(location: str) -> None:
|
|
client, service = harness(PrincipalRole.OPERATOR)
|
|
body = request_body()
|
|
canary = "sensitive-lab-validation-canary"
|
|
if location == "discriminator":
|
|
body["input"]["capability"] = canary
|
|
elif location == "payload":
|
|
body["input"]["payload"]["input"] = {canary: canary}
|
|
else:
|
|
body[canary] = canary
|
|
response = client.post("/api/v1/lab/preflight", json=body)
|
|
assert response.status_code == 422
|
|
assert canary not in response.text
|
|
assert "private-input-canary" not in response.text
|
|
assert response.json()["error"]["code"] == "request_validation_failed"
|
|
service.preflight.assert_not_called()
|