feat(mcp-docugen): Task 4-6 template_store, llm_client, generation_store
- TemplateStore: CRUD filesystem + asset dir, frontmatter YAML roundtrip, path traversal rejection - OpenRouterClient: async httpx con retry backoff esponenziale (5xx, 429, timeout), no-retry su 4xx, parse usage/cost - GenerationStore: SQLite aiosqlite con schema generations + ephemeral_assets, cleanup TTL, stats aggregate Root pyproject aggiornato con respx + pytest-cov dev deps. 19 + 11 + 9 + 6 = 45 test totali, tutti passed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,115 @@
|
||||
import pytest
|
||||
|
||||
from mcp_docugen.generation_store import (
|
||||
EphemeralAssetRecord,
|
||||
GenerationRecord,
|
||||
GenerationStore,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def store(tmp_path):
|
||||
s = GenerationStore(
|
||||
db_path=tmp_path / "gen.db", generated_dir=tmp_path / "generated"
|
||||
)
|
||||
await s.init()
|
||||
return s
|
||||
|
||||
|
||||
async def test_record_success_generation(store):
|
||||
await store.record_generation(
|
||||
GenerationRecord(
|
||||
id="g-1",
|
||||
template_name="fattura",
|
||||
model="m",
|
||||
tokens_in=100,
|
||||
tokens_out=200,
|
||||
cost_usd=0.01,
|
||||
success=True,
|
||||
error_msg=None,
|
||||
)
|
||||
)
|
||||
stats = await store.get_stats()
|
||||
assert stats["total"] == 1
|
||||
assert stats["success"] == 1
|
||||
assert stats["failed"] == 0
|
||||
|
||||
|
||||
async def test_record_failed_generation(store):
|
||||
await store.record_generation(
|
||||
GenerationRecord(
|
||||
id="g-2",
|
||||
template_name="fattura",
|
||||
model="m",
|
||||
tokens_in=0,
|
||||
tokens_out=0,
|
||||
cost_usd=0.0,
|
||||
success=False,
|
||||
error_msg="LLMTimeout",
|
||||
)
|
||||
)
|
||||
stats = await store.get_stats()
|
||||
assert stats["failed"] == 1
|
||||
|
||||
|
||||
async def test_register_ephemeral_asset(store, tmp_path):
|
||||
asset_file = tmp_path / "generated" / "g-1" / "foto.png"
|
||||
asset_file.parent.mkdir(parents=True)
|
||||
asset_file.write_bytes(b"png-bytes")
|
||||
|
||||
await store.register_ephemeral_asset(
|
||||
EphemeralAssetRecord(
|
||||
generation_id="g-1",
|
||||
var_name="foto",
|
||||
file_path=str(asset_file),
|
||||
mime="image/png",
|
||||
ttl_days=30,
|
||||
)
|
||||
)
|
||||
asset = await store.get_ephemeral_asset("g-1", "foto.png")
|
||||
assert asset is not None
|
||||
assert asset.mime == "image/png"
|
||||
|
||||
|
||||
async def test_get_ephemeral_asset_returns_none_if_missing(store):
|
||||
asset = await store.get_ephemeral_asset("nope", "foo.png")
|
||||
assert asset is None
|
||||
|
||||
|
||||
async def test_cleanup_expired_removes_records_and_files(store, tmp_path):
|
||||
asset_file = tmp_path / "generated" / "g-old" / "foto.png"
|
||||
asset_file.parent.mkdir(parents=True)
|
||||
asset_file.write_bytes(b"bytes")
|
||||
|
||||
await store.register_ephemeral_asset(
|
||||
EphemeralAssetRecord(
|
||||
generation_id="g-old",
|
||||
var_name="foto",
|
||||
file_path=str(asset_file),
|
||||
mime="image/png",
|
||||
ttl_days=-1,
|
||||
)
|
||||
)
|
||||
removed = await store.cleanup_expired()
|
||||
assert removed == 1
|
||||
assert not asset_file.exists()
|
||||
assert await store.get_ephemeral_asset("g-old", "foto.png") is None
|
||||
|
||||
|
||||
async def test_ephemeral_asset_expired_flag(store, tmp_path):
|
||||
f = tmp_path / "generated" / "g-e" / "img.png"
|
||||
f.parent.mkdir(parents=True)
|
||||
f.write_bytes(b"x")
|
||||
|
||||
await store.register_ephemeral_asset(
|
||||
EphemeralAssetRecord(
|
||||
generation_id="g-e",
|
||||
var_name="img",
|
||||
file_path=str(f),
|
||||
mime="image/png",
|
||||
ttl_days=-1,
|
||||
)
|
||||
)
|
||||
asset = await store.get_ephemeral_asset("g-e", "img.png")
|
||||
assert asset is not None
|
||||
assert asset.is_expired is True
|
||||
@@ -0,0 +1,173 @@
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
from mcp_docugen.llm_client import (
|
||||
LLMAuthError,
|
||||
LLMEmptyResponse,
|
||||
LLMInvalidResponse,
|
||||
LLMRateLimit,
|
||||
LLMTimeout,
|
||||
LLMUpstreamError,
|
||||
OpenRouterClient,
|
||||
)
|
||||
|
||||
|
||||
def _success_body(text: str = "output text") -> dict:
|
||||
return {
|
||||
"id": "gen-1",
|
||||
"choices": [{"message": {"role": "assistant", "content": text}}],
|
||||
"model": "anthropic/claude-sonnet-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 200,
|
||||
"total_cost": 0.01,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_chat_success():
|
||||
respx.post("https://openrouter.ai/api/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(200, json=_success_body("hello"))
|
||||
)
|
||||
client = OpenRouterClient(
|
||||
api_key="sk", base_url="https://openrouter.ai/api/v1", timeout=5
|
||||
)
|
||||
resp = await client.chat(
|
||||
model="anthropic/claude-sonnet-4", system="sys", user="user"
|
||||
)
|
||||
assert resp.text == "hello"
|
||||
assert resp.tokens_in == 100
|
||||
assert resp.tokens_out == 200
|
||||
assert resp.cost_usd == 0.01
|
||||
assert resp.model == "anthropic/claude-sonnet-4"
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_chat_retries_on_5xx():
|
||||
route = respx.post("https://openrouter.ai/api/v1/chat/completions").mock(
|
||||
side_effect=[
|
||||
httpx.Response(503),
|
||||
httpx.Response(502),
|
||||
httpx.Response(200, json=_success_body()),
|
||||
]
|
||||
)
|
||||
client = OpenRouterClient(
|
||||
api_key="sk",
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
timeout=5,
|
||||
retry_base_delay=0,
|
||||
)
|
||||
resp = await client.chat(model="m", system="s", user="u")
|
||||
assert resp.text == "output text"
|
||||
assert route.call_count == 3
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_chat_exhausts_retries_5xx():
|
||||
respx.post("https://openrouter.ai/api/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(500)
|
||||
)
|
||||
client = OpenRouterClient(
|
||||
api_key="sk",
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
timeout=5,
|
||||
retry_base_delay=0,
|
||||
)
|
||||
with pytest.raises(LLMUpstreamError):
|
||||
await client.chat(model="m", system="s", user="u")
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_chat_retries_on_429():
|
||||
route = respx.post("https://openrouter.ai/api/v1/chat/completions").mock(
|
||||
side_effect=[
|
||||
httpx.Response(429),
|
||||
httpx.Response(200, json=_success_body()),
|
||||
]
|
||||
)
|
||||
client = OpenRouterClient(
|
||||
api_key="sk",
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
timeout=5,
|
||||
retry_base_delay=0,
|
||||
)
|
||||
resp = await client.chat(model="m", system="s", user="u")
|
||||
assert route.call_count == 2
|
||||
assert resp.text == "output text"
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_chat_exhausts_retries_429():
|
||||
respx.post("https://openrouter.ai/api/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(429)
|
||||
)
|
||||
client = OpenRouterClient(
|
||||
api_key="sk",
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
timeout=5,
|
||||
retry_base_delay=0,
|
||||
)
|
||||
with pytest.raises(LLMRateLimit):
|
||||
await client.chat(model="m", system="s", user="u")
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_chat_no_retry_on_401():
|
||||
respx.post("https://openrouter.ai/api/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(401)
|
||||
)
|
||||
client = OpenRouterClient(
|
||||
api_key="sk",
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
timeout=5,
|
||||
retry_base_delay=0,
|
||||
)
|
||||
with pytest.raises(LLMAuthError):
|
||||
await client.chat(model="m", system="s", user="u")
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_chat_timeout():
|
||||
respx.post("https://openrouter.ai/api/v1/chat/completions").mock(
|
||||
side_effect=httpx.ReadTimeout("timeout")
|
||||
)
|
||||
client = OpenRouterClient(
|
||||
api_key="sk",
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
timeout=1,
|
||||
retry_base_delay=0,
|
||||
)
|
||||
with pytest.raises(LLMTimeout):
|
||||
await client.chat(model="m", system="s", user="u")
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_chat_invalid_response_shape():
|
||||
respx.post("https://openrouter.ai/api/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(200, json={"no": "choices"})
|
||||
)
|
||||
client = OpenRouterClient(
|
||||
api_key="sk",
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
timeout=5,
|
||||
retry_base_delay=0,
|
||||
)
|
||||
with pytest.raises(LLMInvalidResponse):
|
||||
await client.chat(model="m", system="s", user="u")
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_chat_empty_content():
|
||||
respx.post("https://openrouter.ai/api/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(200, json=_success_body(text=""))
|
||||
)
|
||||
client = OpenRouterClient(
|
||||
api_key="sk",
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
timeout=5,
|
||||
retry_base_delay=0,
|
||||
)
|
||||
with pytest.raises(LLMEmptyResponse):
|
||||
await client.chat(model="m", system="s", user="u")
|
||||
@@ -0,0 +1,109 @@
|
||||
import base64
|
||||
|
||||
import pytest
|
||||
|
||||
from mcp_docugen.models import TemplateFrontmatter
|
||||
from mcp_docugen.template_store import (
|
||||
InvalidFrontmatter,
|
||||
TemplateAlreadyExists,
|
||||
TemplateNotFound,
|
||||
TemplateStore,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(tmp_path):
|
||||
return TemplateStore(base_dir=tmp_path)
|
||||
|
||||
|
||||
def _fm(name="fattura"):
|
||||
return TemplateFrontmatter(
|
||||
name=name,
|
||||
description="Fattura commerciale",
|
||||
required_variables=[{"name": "cliente", "type": "string"}],
|
||||
)
|
||||
|
||||
|
||||
async def test_create_and_get(store):
|
||||
fm = _fm()
|
||||
await store.create(name="fattura", frontmatter=fm, body="# Body {{cliente}}")
|
||||
got = await store.get("fattura")
|
||||
assert got.frontmatter.name == "fattura"
|
||||
assert "# Body {{cliente}}" in got.body
|
||||
|
||||
|
||||
async def test_create_duplicate_raises(store):
|
||||
fm = _fm()
|
||||
await store.create(name="fattura", frontmatter=fm, body="body")
|
||||
with pytest.raises(TemplateAlreadyExists):
|
||||
await store.create(name="fattura", frontmatter=fm, body="body")
|
||||
|
||||
|
||||
async def test_get_missing_raises(store):
|
||||
with pytest.raises(TemplateNotFound):
|
||||
await store.get("nope")
|
||||
|
||||
|
||||
async def test_list_returns_summaries(store):
|
||||
await store.create(name="a", frontmatter=_fm("a"), body="b")
|
||||
await store.create(name="b", frontmatter=_fm("b"), body="b")
|
||||
result = await store.list()
|
||||
names = sorted(t.name for t in result)
|
||||
assert names == ["a", "b"]
|
||||
|
||||
|
||||
async def test_update_overwrites(store):
|
||||
await store.create(name="f", frontmatter=_fm("f"), body="old")
|
||||
new_fm = TemplateFrontmatter(name="f", description="new description")
|
||||
await store.update(name="f", frontmatter=new_fm, body="new body")
|
||||
got = await store.get("f")
|
||||
assert got.body == "new body"
|
||||
assert got.frontmatter.description == "new description"
|
||||
|
||||
|
||||
async def test_update_missing_raises(store):
|
||||
with pytest.raises(TemplateNotFound):
|
||||
await store.update(name="nope", frontmatter=_fm("nope"), body="x")
|
||||
|
||||
|
||||
async def test_delete_removes(store):
|
||||
await store.create(name="f", frontmatter=_fm("f"), body="b")
|
||||
await store.delete("f")
|
||||
with pytest.raises(TemplateNotFound):
|
||||
await store.get("f")
|
||||
|
||||
|
||||
async def test_delete_missing_raises(store):
|
||||
with pytest.raises(TemplateNotFound):
|
||||
await store.delete("nope")
|
||||
|
||||
|
||||
async def test_assets_are_saved_and_listed(store):
|
||||
png_bytes = b"\x89PNG\r\n\x1a\n"
|
||||
assets = [
|
||||
{
|
||||
"filename": "logo.png",
|
||||
"data_b64": base64.b64encode(png_bytes).decode(),
|
||||
"mime": "image/png",
|
||||
}
|
||||
]
|
||||
await store.create(name="f", frontmatter=_fm("f"), body="b", assets=assets)
|
||||
got = await store.get("f")
|
||||
asset_names = [a.filename for a in got.assets]
|
||||
assert "logo.png" in asset_names
|
||||
content = await store.read_asset("f", "logo.png")
|
||||
assert content == png_bytes
|
||||
|
||||
|
||||
async def test_asset_filename_rejects_path_traversal(store):
|
||||
assets = [{"filename": "../evil.png", "data_b64": "aGk=", "mime": "image/png"}]
|
||||
with pytest.raises(ValueError):
|
||||
await store.create(name="f", frontmatter=_fm("f"), body="b", assets=assets)
|
||||
|
||||
|
||||
async def test_frontmatter_parsing_rejects_malformed_yaml(store, tmp_path):
|
||||
template_dir = tmp_path / "broken"
|
||||
template_dir.mkdir()
|
||||
(template_dir / "template.md").write_text("---\nname: :::broken\n---\nbody")
|
||||
with pytest.raises(InvalidFrontmatter):
|
||||
await store.get("broken")
|
||||
Reference in New Issue
Block a user