test(writer): cover async engine path in WriterAgent/orchestrator
This commit is contained in:
@@ -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"
|
||||||
@@ -47,3 +47,25 @@ def test_generate_produces_filled_docx(tmp_path):
|
|||||||
loaded = Document(str(out))
|
loaded = Document(str(out))
|
||||||
joined = "\n".join(p.text for p in loaded.paragraphs)
|
joined = "\n".join(p.text for p in loaded.paragraphs)
|
||||||
assert "自动生成的内容" in joined
|
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
|
||||||
|
|||||||
Reference in New Issue
Block a user