102 lines
3.5 KiB
Python
102 lines
3.5 KiB
Python
|
"""
|
|||
|
|
Tests for async_execute_prompt (FR-3).
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import pytest
|
|||
|
|
|
|||
|
|
from llm_connect.models import RunConfig, BudgetTracker
|
|||
|
|
from llm_connect.adapter import MockLLMAdapter
|
|||
|
|
from llm_connect.exceptions import LLMBudgetExceededError
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TestAsyncExecutePrompt:
|
|||
|
|
def test_default_fallback_returns_response(self):
|
|||
|
|
adapter = MockLLMAdapter(mock_response="async result")
|
|||
|
|
config = RunConfig()
|
|||
|
|
response = asyncio.run(adapter.async_execute_prompt("hello", config))
|
|||
|
|
assert response.content == "async result"
|
|||
|
|
|
|||
|
|
def test_gather_multiple_adapters(self):
|
|||
|
|
"""asyncio.gather over N adapters completes without errors."""
|
|||
|
|
adapters = [MockLLMAdapter(mock_response=f"resp-{i}") for i in range(4)]
|
|||
|
|
config = RunConfig()
|
|||
|
|
|
|||
|
|
async def run():
|
|||
|
|
return await asyncio.gather(*[
|
|||
|
|
a.async_execute_prompt("prompt", config) for a in adapters
|
|||
|
|
])
|
|||
|
|
|
|||
|
|
results = asyncio.run(run())
|
|||
|
|
assert len(results) == 4
|
|||
|
|
for i, r in enumerate(results):
|
|||
|
|
assert r.content == f"resp-{i}"
|
|||
|
|
|
|||
|
|
def test_gather_increments_call_counts(self):
|
|||
|
|
adapter = MockLLMAdapter()
|
|||
|
|
config = RunConfig()
|
|||
|
|
|
|||
|
|
async def run():
|
|||
|
|
await asyncio.gather(*[
|
|||
|
|
adapter.async_execute_prompt("p", config) for _ in range(5)
|
|||
|
|
])
|
|||
|
|
|
|||
|
|
asyncio.run(run())
|
|||
|
|
assert adapter.call_count == 5
|
|||
|
|
|
|||
|
|
def test_concurrent_faster_than_sequential(self):
|
|||
|
|
"""Gathering N async calls should not be N× slower than one call."""
|
|||
|
|
import time
|
|||
|
|
|
|||
|
|
adapter = MockLLMAdapter()
|
|||
|
|
config = RunConfig()
|
|||
|
|
|
|||
|
|
async def run_concurrent(n: int):
|
|||
|
|
await asyncio.gather(*[
|
|||
|
|
adapter.async_execute_prompt("p", config) for _ in range(n)
|
|||
|
|
])
|
|||
|
|
|
|||
|
|
# Just verify it completes without deadlock or error — timing is CI-unreliable
|
|||
|
|
asyncio.run(run_concurrent(10))
|
|||
|
|
assert adapter.call_count == 10
|
|||
|
|
|
|||
|
|
def test_async_with_budget_tracker(self):
|
|||
|
|
"""Budget enforcement works through async calls."""
|
|||
|
|
tracker = BudgetTracker(total=10000)
|
|||
|
|
config = RunConfig(budget_tracker=tracker)
|
|||
|
|
adapter = MockLLMAdapter(mock_response="hi")
|
|||
|
|
|
|||
|
|
asyncio.run(adapter.async_execute_prompt("hello", config))
|
|||
|
|
assert tracker.spent > 0
|
|||
|
|
|
|||
|
|
def test_async_exhausted_budget_raises(self):
|
|||
|
|
"""Exhausted budget raises LLMBudgetExceededError in async context."""
|
|||
|
|
tracker = BudgetTracker(total=1)
|
|||
|
|
tracker.consume(1)
|
|||
|
|
config = RunConfig(budget_tracker=tracker)
|
|||
|
|
adapter = MockLLMAdapter()
|
|||
|
|
|
|||
|
|
with pytest.raises(LLMBudgetExceededError):
|
|||
|
|
asyncio.run(adapter.async_execute_prompt("p", config))
|
|||
|
|
|
|||
|
|
def test_async_gather_with_shared_budget(self):
|
|||
|
|
"""Shared budget across concurrent async calls is enforced correctly."""
|
|||
|
|
tracker = BudgetTracker(total=100000)
|
|||
|
|
config = RunConfig(budget_tracker=tracker)
|
|||
|
|
adapters = [MockLLMAdapter(mock_response="hi") for _ in range(4)]
|
|||
|
|
|
|||
|
|
async def run():
|
|||
|
|
await asyncio.gather(*[
|
|||
|
|
a.async_execute_prompt("hello", config) for a in adapters
|
|||
|
|
])
|
|||
|
|
|
|||
|
|
asyncio.run(run())
|
|||
|
|
assert tracker.spent > 0
|
|||
|
|
|
|||
|
|
def test_returns_llm_response_type(self):
|
|||
|
|
from llm_connect.models import LLMResponse
|
|||
|
|
adapter = MockLLMAdapter()
|
|||
|
|
config = RunConfig()
|
|||
|
|
response = asyncio.run(adapter.async_execute_prompt("q", config))
|
|||
|
|
assert isinstance(response, LLMResponse)
|