| import threading |
| import time |
|
|
| import pytest |
| from fastapi import HTTPException |
| from fastapi.testclient import TestClient |
|
|
| from server import GenerationRequest, InferenceEngine, create_app, normalize_input |
|
|
|
|
| def test_protein_validation_preserves_one_sequence_and_reports_normalization(): |
| request = GenerationRequest(direction='p2t', input='>candidate\nmaK uZob\nACD') |
| assert request.greedy is False |
| assert request.temperature == 1 |
| assert request.max_length == 512 |
| normalized, note = normalize_input(request) |
| assert normalized == 'MAKXXXXACD' |
| assert note |
| for sequence in ['>one\nMAK\n>two\nMAK', 'MAK123', 'A' * 512, '>header']: |
| with pytest.raises(HTTPException) as error: |
| normalize_input(GenerationRequest(direction='p2t', input=sequence)) |
| assert error.value.status_code == 422 |
|
|
|
|
| def test_api_rejects_invalid_requests_and_does_not_fake_preview_results(): |
| with TestClient(create_app(InferenceEngine(disabled=True))) as client: |
| assert client.get('/api/health').json()['inference'] == 'disabled' |
| assert client.post('/api/jobs', json={'direction': 't2p', 'input': 'PET hydrolase'}).status_code == 503 |
| for payload in [ |
| {'direction': 'invalid', 'input': 'abc'}, |
| {'direction': 't2p', 'input': ' '}, |
| {'direction': 't2p', 'input': 'x', 'max_length': 10000}, |
| {'direction': 'p2t', 'input': 'MAK9'}, |
| ]: |
| assert client.post('/api/jobs', json=payload).status_code == 422 |
| assert client.get('/api/jobs/nonexistent').status_code == 404 |
|
|
|
|
| def test_queue_is_bounded(): |
| engine = InferenceEngine(disabled=False) |
| for _ in range(5): |
| engine.submit(GenerationRequest(direction='t2p', input='Hydrolase')) |
| with pytest.raises(HTTPException) as error: |
| engine.submit(GenerationRequest(direction='t2p', input='Hydrolase')) |
| assert error.value.status_code == 429 |
| assert len(engine.jobs) == 5 |
|
|
|
|
| class RecordingRunner: |
| def __init__(self): |
| self.loads = 0 |
| self.active = 0 |
| self.max_active = 0 |
| self.calls = [] |
|
|
| def load(self): |
| self.loads += 1 |
|
|
| def generate(self, request, sequence): |
| self.active += 1 |
| self.max_active = max(self.active, self.max_active) |
| self.calls.append((request.direction, sequence)) |
| time.sleep(.03) |
| self.active -= 1 |
| if sequence == 'fail': |
| raise ValueError('Please shorten the description.') |
| return 'MAK' if request.direction == 't2p' else 'A test annotation.' |
|
|
|
|
| def test_jobs_run_serially_reuse_model_and_preserve_input_provenance(): |
| runner = RecordingRunner() |
| with TestClient(create_app(InferenceEngine(runner, disabled=False))) as client: |
| first = client.post('/api/jobs', json={'direction': 't2p', 'input': 'Hydrolase'}).json()['id'] |
| second = client.post('/api/jobs', json={'direction': 'p2t', 'input': '>candidate\nma u'}).json()['id'] |
| deadline = time.monotonic() + 3 |
| while time.monotonic() < deadline: |
| b = client.get('/api/jobs/' + second).json() |
| if b['status'] == 'completed': |
| break |
| time.sleep(.02) |
| a = client.get('/api/jobs/' + first).json() |
| assert a['status'] == b['status'] == 'completed' |
| assert a['output'] == 'MAK' |
| assert b['normalized_input'] == 'MAX' |
| assert b['normalization_note'] |
| assert runner.loads == 1 and runner.max_active == 1 |
| assert runner.calls == [('t2p', 'Hydrolase'), ('p2t', 'MAX')] |
|
|
|
|
| def test_failed_generation_does_not_become_a_success_or_poison_the_model(): |
| runner = RecordingRunner() |
| engine = InferenceEngine(runner, disabled=False) |
| engine.start() |
| try: |
| first = engine.submit(GenerationRequest(direction='t2p', input='fail'))['id'] |
| second = engine.submit(GenerationRequest(direction='t2p', input='hydrolase'))['id'] |
| deadline = time.monotonic() + 3 |
| while time.monotonic() < deadline and engine.read(second)['status'] != 'completed': |
| time.sleep(.02) |
| assert engine.read(first)['status'] == 'failed' |
| assert engine.read(first)['output'] is None |
| assert engine.read(second)['status'] == 'completed' |
| assert runner.loads == 1 |
| finally: |
| engine.close() |
|
|
