64 lines
2.5 KiB
Python
64 lines
2.5 KiB
Python
"""Queue lifecycle SQL contracts; real races are covered by the opt-in PG suite."""
|
|
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
|
|
import pytest
|
|
|
|
from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
|
|
from app.mcp.errors import CalculationCapacityExceeded
|
|
|
|
|
|
class Connection:
|
|
def __init__(self, rows=()):
|
|
self.rows = iter(rows)
|
|
self.statements = []
|
|
self.in_transaction = False
|
|
|
|
@asynccontextmanager
|
|
async def transaction(self):
|
|
self.in_transaction = True
|
|
try:
|
|
yield
|
|
finally:
|
|
self.in_transaction = False
|
|
|
|
async def execute(self, sql, params):
|
|
self.statements.append((sql, params, self.in_transaction))
|
|
row = next(self.rows, None) if sql.lstrip().startswith(('SELECT job_id', 'SELECT count', 'INSERT')) else None
|
|
class Cursor:
|
|
async def fetchone(self):
|
|
return row
|
|
async def fetchall(self):
|
|
return []
|
|
return Cursor()
|
|
|
|
|
|
def test_status_projection_excludes_result_and_expires_without_worker():
|
|
connection = Connection()
|
|
repo = CalculationJobRepository(connection, 'project')
|
|
asyncio.run(repo.get('job'))
|
|
assert 'QUEUE_TIMEOUT' in connection.statements[0][0]
|
|
assert 'SET result=NULL' in connection.statements[1][0]
|
|
assert ',result FROM' not in connection.statements[2][0]
|
|
asyncio.run(repo.get('job', include_result=True))
|
|
assert ',result FROM' in connection.statements[-1][0]
|
|
|
|
|
|
def test_duplicate_submission_bypasses_full_queue_inside_transaction():
|
|
existing = {'tool': 'tool', 'parameters': {'timeout_seconds': 3600}, 'binding': {}, 'job_id': 'existing'}
|
|
connection = Connection([existing])
|
|
row = asyncio.run(CalculationJobRepository(connection, 'project').enqueue(
|
|
owner_user_id='owner', request_key='key', tool='tool', parameters={}, binding={}, timeout_seconds=30))
|
|
assert row['job_id'] == 'existing'
|
|
assert 'pg_advisory_xact_lock' in connection.statements[0][0]
|
|
assert all(statement[2] for statement in connection.statements)
|
|
assert not any('count(*)' in sql for sql, _, _ in connection.statements)
|
|
|
|
|
|
def test_capacity_rejects_new_job_before_insert():
|
|
connection = Connection([None, {'active': 100}])
|
|
with pytest.raises(CalculationCapacityExceeded):
|
|
asyncio.run(CalculationJobRepository(connection, 'project').enqueue(
|
|
owner_user_id='owner', request_key='key', tool='tool', parameters={}, binding={}, timeout_seconds=30))
|
|
assert not any('INSERT' in sql for sql, _, _ in connection.statements)
|