| import base64 |
| from io import BytesIO |
| from unittest.mock import patch |
|
|
| import pytest |
| from fastapi.testclient import TestClient |
| from PIL import Image |
| from flux3_engine import Engine, GenerationRequest, LegacyRequest, Settings |
| from web import create_app |
|
|
|
|
| class FakeEngine: |
| def __init__(self): |
| self.pipe = True |
| self.requests = [] |
|
|
| def health(self): |
| return {'ready': True, 'quantization': 'bf16', 'default_steps': 16} |
|
|
| def generate(self, request, image_format='PNG'): |
| self.requests.append(request) |
| data = BytesIO() |
| Image.new('RGB', (request.width, request.height)).save(data, format=image_format) |
| return {'image': base64.b64encode(data.getvalue()).decode(), 'metadata': {}, |
| 'mime_type': 'image/jpeg' if image_format == 'JPEG' else 'image/png'} |
|
|
| legacy = Engine.legacy |
|
|
|
|
| @pytest.fixture |
| def api(monkeypatch): |
| monkeypatch.setenv('FLUX_API_TOKEN', 'test-only') |
| engine = FakeEngine() |
| with TestClient(create_app(engine)) as client: |
| yield client, engine |
|
|
|
|
| def test_auth_required(api): |
| client, engine = api |
| for route in ('/sdxl', '/generate'): |
| assert client.post(route, json={'inputs': 'portrait'}).status_code == 401 |
| assert client.get('/health').status_code == 401 |
| assert not engine.requests |
|
|
|
|
| @pytest.mark.parametrize('route', ['/docs', '/docs/oauth2-redirect', '/redoc', '/openapi.json']) |
| def test_documentation_disabled(api, route): |
| client, engine = api |
| assert client.get(route).status_code == 404 |
| assert client.get(route, headers={'Authorization': 'Bearer test-only'}).status_code == 404 |
| assert not engine.requests |
|
|
|
|
| @pytest.mark.parametrize('startup_error', [None, 'Model load failed']) |
| def test_unready_engine_returns_503(monkeypatch, startup_error): |
| monkeypatch.setenv('FLUX_API_TOKEN', 'test-only') |
| engine = FakeEngine() |
| engine.pipe = True if startup_error else None |
| with TestClient(create_app(engine, startup_error)) as client: |
| headers = {'Authorization': 'Bearer test-only'} |
| assert client.get('/health', headers=headers).status_code == 503 |
| for route in ('/sdxl', '/generate'): |
| assert client.post(route, headers=headers, json={'inputs': 'portrait'}).status_code == 503 |
| assert not engine.requests |
|
|
|
|
| def test_legacy_payload_returns_jpeg_at_16_steps(api): |
| client, engine = api |
| response = client.post('/sdxl', headers={'Authorization': 'Bearer test-only'}, json={ |
| 'inputs': 'portrait', 'num_inference_steps': 40, 'guidance_scale': 6, |
| 'negative_prompt': 'blurry', 'use_refiner': False, 'high_noise_frac': .8, |
| 'width': 512, 'height': 768}) |
| assert response.status_code == 200 |
| image = Image.open(BytesIO(base64.b64decode(response.json()['image']))) |
| assert image.format == 'JPEG' and image.size == (512, 768) |
| assert engine.requests[0].num_inference_steps == 16 |
|
|
|
|
| def test_native_resolution_and_steps(api): |
| client, engine = api |
| headers = {'Authorization': 'Bearer test-only'} |
| assert client.post('/generate', headers=headers, |
| json={'inputs': 'portrait', 'width': 768, 'height': 768}).status_code == 200 |
| assert engine.requests[-1].num_inference_steps == 16 |
| for payload in ({'inputs': ' '}, {'inputs': 'p', 'width': 1000}, |
| {'inputs': 'p', 'num_inference_steps': 40}): |
| assert client.post('/generate', headers=headers, json=payload).status_code == 422 |
|
|
|
|
| def test_fail_closed_without_api_token(monkeypatch): |
| monkeypatch.delenv('FLUX_API_TOKEN', raising=False) |
| with pytest.raises(RuntimeError): |
| create_app(FakeEngine()) |
|
|
|
|
| def test_private_repo_does_not_reuse_upstream_commit(monkeypatch): |
| monkeypatch.setenv('FLUX_MODEL_ID', 'example/private-model') |
| monkeypatch.delenv('FLUX_MODEL_REVISION', raising=False) |
| monkeypatch.delenv('FLUX_QUANTIZATION', raising=False) |
| assert Settings.from_env() == Settings('example/private-model', 'main', 'q8_0') |
|
|
|
|
| def test_hf_handler_uses_mounted_repository(): |
| with patch('handler.Engine') as engine: |
| from handler import EndpointHandler |
| handler = EndpointHandler('/repository') |
| payload = {'inputs': 'portrait'} |
| handler(payload) |
| engine.assert_called_once_with(model_path='/repository') |
| engine.return_value.load.assert_called_once() |
| engine.return_value.legacy.assert_called_once_with(payload) |
|
|
