749 lines
26 KiB
Python
749 lines
26 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import pytest
|
|
from sqlalchemy import create_engine, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
|
|
from app.core.db import Base
|
|
from app.core.task_manager import DeliveryMode, SqlAlchemyTaskStore
|
|
from app.core.task_manager.types import TaskStatus
|
|
from app.models.task import GenerationTask, GenerationTaskStatus
|
|
from app.models.task_links import GenerationTaskLink
|
|
from app.services.script_processing_worker import run_extract_task_sync
|
|
from app.services.script_processing_tasks import (
|
|
CHARACTER_PORTRAIT_ANALYSIS_RELATION_TYPE,
|
|
CHAPTER_DIVISION_RELATION_TYPE,
|
|
CONSISTENCY_CHECK_RELATION_TYPE,
|
|
COSTUME_INFO_ANALYSIS_RELATION_TYPE,
|
|
ENTITY_MERGE_RELATION_TYPE,
|
|
PROP_INFO_ANALYSIS_RELATION_TYPE,
|
|
SCENE_INFO_ANALYSIS_RELATION_TYPE,
|
|
SCRIPT_OPTIMIZATION_RELATION_TYPE,
|
|
SCRIPT_SIMPLIFICATION_RELATION_TYPE,
|
|
SCRIPT_EXTRACTION_RELATION_TYPE,
|
|
VARIANT_ANALYSIS_RELATION_TYPE,
|
|
create_character_portrait_task,
|
|
create_consistency_task,
|
|
create_costume_info_task,
|
|
create_divide_task,
|
|
create_extract_task,
|
|
create_merge_task,
|
|
create_prop_info_task,
|
|
create_scene_info_task,
|
|
create_script_optimization_task,
|
|
create_script_simplification_task,
|
|
create_variant_task,
|
|
pick_analysis_relation_entity_id,
|
|
pick_consistency_relation_entity_id,
|
|
pick_merge_relation_entity_id,
|
|
pick_variant_relation_entity_id,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_divide_task_creates_task_and_link() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
result = await create_divide_task(
|
|
db,
|
|
chapter_id="chapter-1",
|
|
script_text="一段剧本",
|
|
write_to_db=True,
|
|
)
|
|
|
|
assert result.reused is False
|
|
assert result.status == TaskStatus.pending
|
|
assert result.relation_type == CHAPTER_DIVISION_RELATION_TYPE
|
|
assert result.relation_entity_id == "chapter-1"
|
|
|
|
task = await db.get(GenerationTask, result.task_id)
|
|
assert task is not None
|
|
|
|
link = (
|
|
await db.execute(
|
|
select(GenerationTaskLink).where(
|
|
GenerationTaskLink.task_id == result.task_id,
|
|
GenerationTaskLink.relation_type == CHAPTER_DIVISION_RELATION_TYPE,
|
|
GenerationTaskLink.relation_entity_id == "chapter-1",
|
|
)
|
|
)
|
|
).scalars().first()
|
|
assert link is not None
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_divide_task_reuses_existing_active_task() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
first = await create_divide_task(
|
|
db,
|
|
chapter_id="chapter-1",
|
|
script_text="第一版",
|
|
write_to_db=True,
|
|
)
|
|
second = await create_divide_task(
|
|
db,
|
|
chapter_id="chapter-1",
|
|
script_text="第二版",
|
|
write_to_db=True,
|
|
)
|
|
|
|
assert second.reused is True
|
|
assert second.task_id == first.task_id
|
|
|
|
link_count = len(
|
|
(
|
|
await db.execute(
|
|
select(GenerationTaskLink).where(
|
|
GenerationTaskLink.relation_type == CHAPTER_DIVISION_RELATION_TYPE,
|
|
GenerationTaskLink.relation_entity_id == "chapter-1",
|
|
)
|
|
)
|
|
)
|
|
.scalars()
|
|
.all()
|
|
)
|
|
assert link_count == 1
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_request_cancel_marks_pending_task_cancelled_immediately() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
store = SqlAlchemyTaskStore(db)
|
|
task = await store.create(payload={"k": "v"}, mode=DeliveryMode.async_polling, task_kind="test_task")
|
|
rec = await store.request_cancel(task.id, "用户取消")
|
|
assert rec is not None
|
|
assert rec.cancel_requested is True
|
|
assert rec.status == TaskStatus.cancelled
|
|
|
|
view = await store.get_status_view(task.id)
|
|
assert view is not None
|
|
assert view.cancel_requested is True
|
|
assert view.status == TaskStatus.cancelled
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_extract_task_creates_task_and_link() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
result = await create_extract_task(
|
|
db,
|
|
project_id="project-1",
|
|
chapter_id="chapter-1",
|
|
script_division={"shots": []},
|
|
consistency=None,
|
|
refresh_cache=False,
|
|
)
|
|
|
|
assert result.reused is False
|
|
assert result.status == TaskStatus.pending
|
|
assert result.relation_type == SCRIPT_EXTRACTION_RELATION_TYPE
|
|
assert result.relation_entity_id == "chapter-1"
|
|
|
|
task = await db.get(GenerationTask, result.task_id)
|
|
assert task is not None
|
|
|
|
link = (
|
|
await db.execute(
|
|
select(GenerationTaskLink).where(
|
|
GenerationTaskLink.task_id == result.task_id,
|
|
GenerationTaskLink.relation_type == SCRIPT_EXTRACTION_RELATION_TYPE,
|
|
GenerationTaskLink.relation_entity_id == "chapter-1",
|
|
)
|
|
)
|
|
).scalars().first()
|
|
assert link is not None
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_extract_task_reuses_existing_active_task() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
first = await create_extract_task(
|
|
db,
|
|
project_id="project-1",
|
|
chapter_id="chapter-1",
|
|
script_division={"shots": []},
|
|
consistency=None,
|
|
refresh_cache=False,
|
|
)
|
|
second = await create_extract_task(
|
|
db,
|
|
project_id="project-1",
|
|
chapter_id="chapter-1",
|
|
script_division={"shots": [{"index": 1}]},
|
|
consistency={"ok": True},
|
|
refresh_cache=True,
|
|
)
|
|
|
|
assert second.reused is True
|
|
assert second.task_id == first.task_id
|
|
|
|
link_count = len(
|
|
(
|
|
await db.execute(
|
|
select(GenerationTaskLink).where(
|
|
GenerationTaskLink.relation_type == SCRIPT_EXTRACTION_RELATION_TYPE,
|
|
GenerationTaskLink.relation_entity_id == "chapter-1",
|
|
)
|
|
)
|
|
)
|
|
.scalars()
|
|
.all()
|
|
)
|
|
assert link_count == 1
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_merge_task_creates_task_and_link() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
result = await create_merge_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
all_shot_extractions=[],
|
|
historical_library=None,
|
|
script_division=None,
|
|
previous_merge=None,
|
|
conflict_resolutions=None,
|
|
)
|
|
|
|
assert result.reused is False
|
|
assert result.status == TaskStatus.pending
|
|
assert result.relation_type == ENTITY_MERGE_RELATION_TYPE
|
|
assert result.relation_entity_id == "chapter-1"
|
|
|
|
link = (
|
|
await db.execute(
|
|
select(GenerationTaskLink).where(
|
|
GenerationTaskLink.task_id == result.task_id,
|
|
GenerationTaskLink.relation_type == ENTITY_MERGE_RELATION_TYPE,
|
|
GenerationTaskLink.relation_entity_id == "chapter-1",
|
|
)
|
|
)
|
|
).scalars().first()
|
|
assert link is not None
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_merge_task_reuses_existing_active_task() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
first = await create_merge_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
all_shot_extractions=[],
|
|
historical_library=None,
|
|
script_division=None,
|
|
previous_merge=None,
|
|
conflict_resolutions=None,
|
|
)
|
|
second = await create_merge_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
all_shot_extractions=[{"a": 1}],
|
|
historical_library={"x": 1},
|
|
script_division={"shots": []},
|
|
previous_merge=None,
|
|
conflict_resolutions=[],
|
|
)
|
|
|
|
assert second.reused is True
|
|
assert second.task_id == first.task_id
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
def test_pick_merge_relation_entity_id_prefers_chapter() -> None:
|
|
assert pick_merge_relation_entity_id(chapter_id="chapter-1", project_id="project-1") == "chapter-1"
|
|
assert pick_merge_relation_entity_id(chapter_id=None, project_id="project-1") == "project-1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_consistency_task_creates_task_and_link() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
result = await create_consistency_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
script_text="完整剧本",
|
|
)
|
|
|
|
assert result.reused is False
|
|
assert result.status == TaskStatus.pending
|
|
assert result.relation_type == CONSISTENCY_CHECK_RELATION_TYPE
|
|
assert result.relation_entity_id == "chapter-1"
|
|
|
|
link = (
|
|
await db.execute(
|
|
select(GenerationTaskLink).where(
|
|
GenerationTaskLink.task_id == result.task_id,
|
|
GenerationTaskLink.relation_type == CONSISTENCY_CHECK_RELATION_TYPE,
|
|
GenerationTaskLink.relation_entity_id == "chapter-1",
|
|
)
|
|
)
|
|
).scalars().first()
|
|
assert link is not None
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_consistency_task_reuses_existing_active_task() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
first = await create_consistency_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
script_text="完整剧本 A",
|
|
)
|
|
second = await create_consistency_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
script_text="完整剧本 B",
|
|
)
|
|
|
|
assert second.reused is True
|
|
assert second.task_id == first.task_id
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
def test_pick_consistency_relation_entity_id_prefers_chapter() -> None:
|
|
assert pick_consistency_relation_entity_id(chapter_id="chapter-1", project_id="project-1") == "chapter-1"
|
|
assert pick_consistency_relation_entity_id(chapter_id=None, project_id="project-1") == "project-1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_variant_task_creates_task_and_link() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
result = await create_variant_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
merged_library={},
|
|
all_shot_extractions=[],
|
|
script_division=None,
|
|
)
|
|
|
|
assert result.reused is False
|
|
assert result.status == TaskStatus.pending
|
|
assert result.relation_type == VARIANT_ANALYSIS_RELATION_TYPE
|
|
assert result.relation_entity_id == "chapter-1"
|
|
|
|
link = (
|
|
await db.execute(
|
|
select(GenerationTaskLink).where(
|
|
GenerationTaskLink.task_id == result.task_id,
|
|
GenerationTaskLink.relation_type == VARIANT_ANALYSIS_RELATION_TYPE,
|
|
GenerationTaskLink.relation_entity_id == "chapter-1",
|
|
)
|
|
)
|
|
).scalars().first()
|
|
assert link is not None
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_variant_task_reuses_existing_active_task() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
first = await create_variant_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
merged_library={},
|
|
all_shot_extractions=[],
|
|
script_division=None,
|
|
)
|
|
second = await create_variant_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
merged_library={"characters": []},
|
|
all_shot_extractions=[{"a": 1}],
|
|
script_division={"shots": []},
|
|
)
|
|
|
|
assert second.reused is True
|
|
assert second.task_id == first.task_id
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
def test_pick_variant_relation_entity_id_prefers_chapter() -> None:
|
|
assert pick_variant_relation_entity_id(chapter_id="chapter-1", project_id="project-1") == "chapter-1"
|
|
assert pick_variant_relation_entity_id(chapter_id=None, project_id="project-1") == "project-1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_character_portrait_task_creates_task_and_link() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
result = await create_character_portrait_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
character_context=None,
|
|
character_description="人物描述",
|
|
)
|
|
assert result.relation_type == CHARACTER_PORTRAIT_ANALYSIS_RELATION_TYPE
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_prop_info_task_creates_task_and_link() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
result = await create_prop_info_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
prop_context=None,
|
|
prop_description="道具描述",
|
|
)
|
|
assert result.relation_type == PROP_INFO_ANALYSIS_RELATION_TYPE
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_scene_info_task_creates_task_and_link() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
result = await create_scene_info_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
scene_context=None,
|
|
scene_description="场景描述",
|
|
)
|
|
assert result.relation_type == SCENE_INFO_ANALYSIS_RELATION_TYPE
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_costume_info_task_creates_task_and_link() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
result = await create_costume_info_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
costume_context=None,
|
|
costume_description="服装描述",
|
|
)
|
|
assert result.relation_type == COSTUME_INFO_ANALYSIS_RELATION_TYPE
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_script_optimization_task_creates_task_and_link() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
result = await create_script_optimization_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
script_text="原始剧本",
|
|
consistency={"has_issues": True},
|
|
)
|
|
assert result.relation_type == SCRIPT_OPTIMIZATION_RELATION_TYPE
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_script_simplification_task_creates_task_and_link() -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
result = await create_script_simplification_task(
|
|
db,
|
|
relation_entity_id="chapter-1",
|
|
script_text="原始剧本",
|
|
)
|
|
assert result.relation_type == SCRIPT_SIMPLIFICATION_RELATION_TYPE
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
def test_pick_analysis_relation_entity_id_prefers_chapter() -> None:
|
|
assert pick_analysis_relation_entity_id(chapter_id="chapter-1", project_id="project-1", endpoint="x") == "chapter-1"
|
|
assert pick_analysis_relation_entity_id(chapter_id=None, project_id="project-1", endpoint="x") == "project-1"
|
|
|
|
|
|
@dataclass
|
|
class _FakeDraft:
|
|
payload: dict
|
|
|
|
def model_dump(self) -> dict:
|
|
return self.payload
|
|
|
|
def model_copy(self, *, deep: bool = False): # noqa: ARG002
|
|
return _FakeDraft(dict(self.payload))
|
|
|
|
|
|
class _FakeExtractorAgent:
|
|
def __init__(self, _llm) -> None:
|
|
pass
|
|
|
|
def extract(self, **_kwargs):
|
|
return _FakeDraft({"shots": []})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_extract_task_marks_cancelled_when_cancel_requested_before_start(monkeypatch, tmp_path) -> None:
|
|
db_path = tmp_path / "extract-cancel-before-start.db"
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}", future=True)
|
|
sync_engine = create_engine(f"sqlite:///{db_path}", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
SyncSessionLocal = sessionmaker(sync_engine, class_=Session, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async with SessionLocal() as db:
|
|
created = await create_extract_task(
|
|
db,
|
|
project_id="project-1",
|
|
chapter_id="chapter-1",
|
|
script_division={"shots": []},
|
|
consistency=None,
|
|
refresh_cache=False,
|
|
)
|
|
store = SqlAlchemyTaskStore(db)
|
|
await store.request_cancel(created.task_id, "用户取消")
|
|
await db.commit()
|
|
|
|
monkeypatch.setattr("app.services.script_processing_worker.sync_session_maker", SyncSessionLocal)
|
|
|
|
run_extract_task_sync(created.task_id)
|
|
|
|
async with SessionLocal() as db:
|
|
row = await db.get(GenerationTask, created.task_id)
|
|
assert row is not None
|
|
assert row.status == GenerationTaskStatus.cancelled
|
|
assert bool(row.cancel_requested) is True
|
|
assert row.cancelled_at is not None
|
|
|
|
await engine.dispose()
|
|
sync_engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_extract_task_marks_cancelled_at_stage_boundary(monkeypatch, tmp_path) -> None:
|
|
db_path = tmp_path / "extract-cancel-stage-boundary.db"
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}", future=True)
|
|
sync_engine = create_engine(f"sqlite:///{db_path}", future=True)
|
|
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
SyncSessionLocal = sessionmaker(sync_engine, class_=Session, expire_on_commit=False)
|
|
|
|
import app.models.task # noqa: F401
|
|
import app.models.task_links # noqa: F401
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
monkeypatch.setattr("app.services.script_processing_worker.sync_session_maker", SyncSessionLocal)
|
|
monkeypatch.setattr("app.services.script_processing_worker.ElementExtractorAgent", _FakeExtractorAgent)
|
|
|
|
def _fake_llm(_db, *, thinking: bool): # noqa: ARG001
|
|
return object()
|
|
|
|
def _cancel_during_sync(db, chapter_id: str, draft) -> None: # noqa: ANN001
|
|
del chapter_id, draft
|
|
from app.core.task_manager import SyncSqlAlchemyTaskStore
|
|
|
|
store = SyncSqlAlchemyTaskStore(db)
|
|
task_id = getattr(_cancel_during_sync, "task_id")
|
|
store.request_cancel(task_id, "阶段边界取消")
|
|
db.commit()
|
|
|
|
def _noop_dialog_sync(_db, chapter_id: str, draft) -> None: # noqa: ANN001
|
|
return None
|
|
|
|
monkeypatch.setattr("app.services.script_processing_worker.build_default_text_llm_sync", _fake_llm)
|
|
monkeypatch.setattr(
|
|
"app.services.script_processing_worker.sync_shot_extracted_candidates_from_draft_sync",
|
|
_cancel_during_sync,
|
|
)
|
|
monkeypatch.setattr(
|
|
"app.services.script_processing_worker.sync_shot_extracted_dialogue_candidates_from_draft_sync",
|
|
_noop_dialog_sync,
|
|
)
|
|
|
|
async with SessionLocal() as db:
|
|
created = await create_extract_task(
|
|
db,
|
|
project_id="project-1",
|
|
chapter_id="chapter-1",
|
|
script_division={"shots": []},
|
|
consistency=None,
|
|
refresh_cache=True,
|
|
)
|
|
await db.commit()
|
|
setattr(_cancel_during_sync, "task_id", created.task_id)
|
|
|
|
run_extract_task_sync(created.task_id)
|
|
|
|
async with SessionLocal() as db:
|
|
row = await db.get(GenerationTask, created.task_id)
|
|
assert row is not None
|
|
assert row.status == GenerationTaskStatus.cancelled
|
|
assert bool(row.cancel_requested) is True
|
|
assert row.result == {"draft": {"shots": []}, "from_cache": False}
|
|
|
|
await engine.dispose()
|
|
sync_engine.dispose()
|