m4
229 lines
7.1 KiB
Python
229 lines
7.1 KiB
Python
"""Contract tests for the /api/image-generation endpoints."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
import jwt
|
|
import pytest
|
|
import yaml
|
|
from cryptography.hazmat.primitives import serialization
|
|
from cryptography.hazmat.primitives.asymmetric import ec
|
|
from starlette.applications import Starlette
|
|
from starlette.testclient import TestClient
|
|
|
|
from EvoScientist.image_gen import config as image_config
|
|
from EvoScientist.image_gen.config import (
|
|
ImageGenerationSettings,
|
|
ImageModelEntry,
|
|
save_image_generation_settings,
|
|
)
|
|
from EvoScientist.model_registry import http_api
|
|
from EvoScientist.model_registry.auth import BffAuthenticator
|
|
from EvoScientist.model_registry.endpoint_policy import EndpointPolicy
|
|
from EvoScientist.model_registry.hashing import configuration_hash # noqa: F401 (fixture parity)
|
|
from EvoScientist.model_registry.http_api import ApiServices, model_registry_routes
|
|
from EvoScientist.model_registry.platform import DelegationPublicKey
|
|
from EvoScientist.model_registry.resolver import ModelRegistryResolver
|
|
from EvoScientist.model_registry.snapshots import SnapshotService
|
|
from EvoScientist.model_registry.store import ModelRuntimeStore
|
|
|
|
SERVICE_TOKEN = "bff-service-token"
|
|
|
|
_PRIVATE_KEY = ec.generate_private_key(ec.SECP256R1())
|
|
_PRIVATE_PEM = _PRIVATE_KEY.private_bytes(
|
|
serialization.Encoding.PEM,
|
|
serialization.PrivateFormat.PKCS8,
|
|
serialization.NoEncryption(),
|
|
)
|
|
_PUBLIC_PEM = _PRIVATE_KEY.public_key().public_bytes(
|
|
serialization.Encoding.PEM,
|
|
serialization.PublicFormat.SubjectPublicKeyInfo,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def config_path(tmp_path):
|
|
path = tmp_path / "config.yaml"
|
|
path.write_text("other_section: {keep: true}\n", encoding="utf-8")
|
|
return path
|
|
|
|
|
|
@pytest.fixture
|
|
def services(tmp_path, config_path, monkeypatch):
|
|
store = ModelRuntimeStore(config_dir=tmp_path / "runtime")
|
|
monkeypatch.setattr(
|
|
http_api,
|
|
"load_image_generation_settings",
|
|
lambda: image_config.load_image_generation_settings(config_path=config_path),
|
|
)
|
|
monkeypatch.setattr(
|
|
http_api,
|
|
"save_image_generation_settings",
|
|
lambda settings: image_config.save_image_generation_settings(
|
|
settings, config_path=config_path
|
|
),
|
|
)
|
|
resolver = ModelRegistryResolver(store)
|
|
return ApiServices(
|
|
store=store,
|
|
resolver=resolver,
|
|
snapshot_service=SnapshotService(store, resolver),
|
|
endpoint_policy=EndpointPolicy([]),
|
|
authenticator=BffAuthenticator(
|
|
service_token=SERVICE_TOKEN,
|
|
service_token_hash=None,
|
|
delegation_keys=(
|
|
DelegationPublicKey(deployment_id="webui-1", public_key=_PUBLIC_PEM),
|
|
),
|
|
jti_store=store,
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def client(services):
|
|
app = Starlette(routes=model_registry_routes(lambda: services))
|
|
return TestClient(app)
|
|
|
|
|
|
def _headers(scopes):
|
|
now = int(time.time())
|
|
claims = {
|
|
"iss": "WebUI",
|
|
"aud": "EvoScientist",
|
|
"sub": "user-1",
|
|
"scopes": list(scopes),
|
|
"deployment_id": "webui-1",
|
|
"iat": now,
|
|
"exp": now + 30,
|
|
"jti": uuid.uuid4().hex,
|
|
}
|
|
return {
|
|
"Authorization": f"Bearer {SERVICE_TOKEN}",
|
|
"X-Evo-Actor": jwt.encode(claims, _PRIVATE_PEM, algorithm="ES256"),
|
|
}
|
|
|
|
|
|
def _admin_headers():
|
|
return _headers(["model_config:read", "model_config:write"])
|
|
|
|
|
|
def _settings_payload(**overrides):
|
|
payload = {
|
|
"default_model": "gpt-image-2",
|
|
"timeout_seconds": 120.0,
|
|
"models": [
|
|
{
|
|
"id": "gpt-image-2",
|
|
"name": "GPT Image 2",
|
|
"provider": "openai",
|
|
"api_key": "sk-live-9876abcd",
|
|
"base_url": "",
|
|
"supports_generation": True,
|
|
"supports_edit": True,
|
|
"default_size": "1024x1024",
|
|
"default_quality": "auto",
|
|
"params": {},
|
|
}
|
|
],
|
|
}
|
|
payload.update(overrides)
|
|
return payload
|
|
|
|
|
|
def test_put_then_get_masks_api_key(client, config_path):
|
|
put = client.put(
|
|
"/api/image-generation", headers=_admin_headers(), json=_settings_payload()
|
|
)
|
|
assert put.status_code == 200
|
|
body = put.json()
|
|
model = body["models"][0]
|
|
assert model["api_key"] == ""
|
|
assert model["api_key_configured"] is True
|
|
assert model["api_key_hint"] == "...abcd"
|
|
assert "sk-live-9876abcd" not in json.dumps(body)
|
|
|
|
get = client.get("/api/image-generation", headers=_admin_headers())
|
|
assert get.status_code == 200
|
|
assert get.json() == body
|
|
|
|
# The plaintext key really landed in config.yaml, other sections intact.
|
|
raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
|
assert raw["other_section"] == {"keep": True}
|
|
assert raw["image_generation"]["models"][0]["api_key"] == "sk-live-9876abcd"
|
|
|
|
|
|
def test_put_blank_api_key_keeps_existing(client):
|
|
client.put(
|
|
"/api/image-generation", headers=_admin_headers(), json=_settings_payload()
|
|
)
|
|
update = _settings_payload()
|
|
update["models"][0]["api_key"] = ""
|
|
update["models"][0]["name"] = "Renamed"
|
|
response = client.put(
|
|
"/api/image-generation", headers=_admin_headers(), json=update
|
|
)
|
|
assert response.status_code == 200
|
|
model = response.json()["models"][0]
|
|
assert model["name"] == "Renamed"
|
|
assert model["api_key_configured"] is True
|
|
assert model["api_key_hint"] == "...abcd"
|
|
|
|
|
|
def test_put_new_model_without_key_is_unconfigured(client):
|
|
payload = _settings_payload(models=[_settings_payload()["models"][0] | {"api_key": ""}])
|
|
response = client.put(
|
|
"/api/image-generation", headers=_admin_headers(), json=payload
|
|
)
|
|
assert response.status_code == 200
|
|
model = response.json()["models"][0]
|
|
assert model["api_key_configured"] is False
|
|
assert model["api_key_hint"] is None
|
|
|
|
|
|
def test_short_key_never_leaks_through_hint(client):
|
|
payload = _settings_payload()
|
|
payload["models"][0]["api_key"] = "abc"
|
|
response = client.put(
|
|
"/api/image-generation", headers=_admin_headers(), json=payload
|
|
)
|
|
assert response.status_code == 200
|
|
model = response.json()["models"][0]
|
|
assert model["api_key_configured"] is True
|
|
assert model["api_key_hint"] is None
|
|
|
|
|
|
def test_get_empty_when_unconfigured(client):
|
|
response = client.get("/api/image-generation", headers=_admin_headers())
|
|
assert response.status_code == 200
|
|
assert response.json() == {
|
|
"default_model": "",
|
|
"timeout_seconds": 120.0,
|
|
"models": [],
|
|
}
|
|
|
|
|
|
def test_put_rejects_invalid_payload(client):
|
|
payload = _settings_payload()
|
|
payload["models"][0]["provider"] = "midjourney"
|
|
response = client.put(
|
|
"/api/image-generation", headers=_admin_headers(), json=payload
|
|
)
|
|
assert response.status_code == 422
|
|
body = response.json()
|
|
assert body["code"] == "VALIDATION_FAILED"
|
|
assert "midjourney" not in json.dumps(body)
|
|
|
|
|
|
def test_read_scope_cannot_write(client):
|
|
response = client.put(
|
|
"/api/image-generation",
|
|
headers=_headers(["model_config:read"]),
|
|
json=_settings_payload(),
|
|
)
|
|
assert response.status_code == 403
|