File size: 4,319 Bytes
03677d0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | 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()
|