test(writer): cover async engine path in WriterAgent/orchestrator

This commit is contained in:
lhl
2026-08-13 11:49:44 +08:00
parent e2b4faedbe
commit 2770c76ba0
2 changed files with 112 additions and 0 deletions
+90
View File
@@ -0,0 +1,90 @@
import pytest
from types import SimpleNamespace
from genesis.writer.writer_agent import (
WriterAgent,
WRITER_PROMPT_TEMPLATE,
CHAPTER_OUTPUT_SCHEMA,
)
from genesis.writer.models import GenerationContext, ChapterSpec
from genesis.writer.writer_state import WriterState
from genesis.writer.exceptions import WriterGenerationError
class FakeEngine:
def __init__(self):
self.calls = 0
def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
self.calls += 1
return SimpleNamespace(
data={"title": variables["title"], "blocks": [{"type": "paragraph", "text": "ok"}]},
status="ok",
)
class FakePromptRegistry:
@staticmethod
def get_or_create(name, template):
return SimpleNamespace(name=name, version="1", template=template)
def _ctx(cid, title):
return GenerationContext(
chapter_id=cid,
title=title,
template_marker=ChapterSpec(chapter_id=cid, title=title, section_placeholder=None),
structured_source="source text",
write_rules=["W1"],
design_rules=["D1"],
template_styles={"Heading1"},
prior_state=None,
)
def test_generate_chapter_success():
agent = WriterAgent(
session_id="s",
engine=FakeEngine(),
prompt_registry=FakePromptRegistry(),
state=WriterState(["db_design"]),
)
content = agent.generate_chapter(_ctx("db_design", "DB 設計"))
assert content.chapter_id == "db_design"
assert content.blocks and content.blocks[0].type == "paragraph"
def test_generate_chapter_retries():
class Boom(FakeEngine):
def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
raise RuntimeError("boom")
agent = WriterAgent(
session_id="s",
engine=Boom(),
prompt_registry=FakePromptRegistry(),
state=WriterState(["db_design"]),
max_retries=2,
)
with pytest.raises(WriterGenerationError):
agent.generate_chapter(_ctx("db_design", "DB 設計"))
class AsyncFakeEngine:
async def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
return SimpleNamespace(
data={"title": variables["title"], "blocks": [{"type": "paragraph", "text": "ok"}]},
status="ok",
)
def test_generate_chapter_with_async_engine():
agent = WriterAgent(
session_id="s",
engine=AsyncFakeEngine(),
prompt_registry=FakePromptRegistry(),
state=WriterState(["db_design"]),
)
content = agent.generate_chapter(_ctx("db_design", "DB 設計"))
assert content.chapter_id == "db_design"
assert content.blocks and content.blocks[0].type == "paragraph"
+22
View File
@@ -47,3 +47,25 @@ def test_generate_produces_filled_docx(tmp_path):
loaded = Document(str(out))
joined = "\n".join(p.text for p in loaded.paragraphs)
assert "自动生成的内容" in joined
class AsyncFakeEngine:
async def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
return SimpleNamespace(
data={"title": variables["title"], "blocks": [{"type": "paragraph", "text": "异步引擎内容"}]},
status="ok",
)
def test_generate_with_async_engine_produces_filled_docx(tmp_path):
tpl = tmp_path / "tpl.docx"
out = tmp_path / "out.docx"
_make_template(str(tpl))
orch = WriteOrchestrator()
contents = orch.generate(
_ss(str(tpl)), str(out), samples_dir="nonexistent_dir_xyz", engine=AsyncFakeEngine()
)
assert isinstance(contents, list) and len(contents) == 1
loaded = Document(str(out))
joined = "\n".join(p.text for p in loaded.paragraphs)
assert "异步引擎内容" in joined