PMFlow/backend/tests/test_mounted_ai_runtime_models.py

from __future__ import annotations

import os

import pytest
from sqlalchemy import JSON, create_engine
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session

os.environ["DATABASE_URL"] = "sqlite:///:memory:"

from app.models.base import Base
from app.models.mounted_ai_runtime import (
    MountedAgentTrace,
    MountedConversationRun,
    MountedDraftWorkspace,
    MountedSourceActivation,
    MountedToolCapabilityPolicy,
    MountedTopicWorkspace,
)


def _use_sqlite_json_types() -> None:
    for table in (
        MountedTopicWorkspace.__table__,
        MountedDraftWorkspace.__table__,
        MountedConversationRun.__table__,
        MountedAgentTrace.__table__,
        MountedSourceActivation.__table__,
        MountedToolCapabilityPolicy.__table__,
    ):
        for column in table.columns:
            if column.name in {
                "facts_json",
                "goals_json",
                "constraints_json",
                "options_json",
                "open_questions",
                "pending_actions",
                "bound_entities",
                "metadata_json",
                "requirement_card",
                "draft_payload",
                "required_confirmations",
                "context_layers_hit",
                "tool_calls",
                "execution_outcome",
                "input_summary",
                "output_summary",
                "tool_trace",
                "source_scope",
                "tenant_scope",
                "policy_metadata",
            }:
                column.type = JSON()


def _create_session() -> Session:
    engine = create_engine("sqlite:///:memory:")
    _use_sqlite_json_types()
    Base.metadata.create_all(
        engine,
        tables=[
            MountedTopicWorkspace.__table__,
            MountedDraftWorkspace.__table__,
            MountedConversationRun.__table__,
            MountedAgentTrace.__table__,
            MountedSourceActivation.__table__,
            MountedToolCapabilityPolicy.__table__,
        ],
    )
    return Session(engine)


def test_mounted_ai_runtime_models_roundtrip_and_defaults() -> None:
    session = _create_session()

    topic = MountedTopicWorkspace(
        tenant_id=1,
        topic_key="feishu:chat_1:topic_1",
        source_system="feishu",
        channel_ref="chat_1",
        thread_key="thread_1",
        root_message_id="om_root_1",
        mode="plan",
        facts_json={"problem": "直播复盘"},
        open_questions=["负责人是谁"],
        readiness_score=0.35,
    )
    session.add(topic)
    session.flush()

    draft = MountedDraftWorkspace(
        tenant_id=1,
        topic_workspace_id=topic.id,
        draft_version=1,
        requirement_card={"goals": ["整理复盘需求"]},
        required_confirmations=[{"type": "user_confirm"}],
        active_work_item_id=501,
        latest_source_message_id="om_reply_source_1",
        latest_reply_message_id="om_reply_sent_1",
    )
    run = MountedConversationRun(
        tenant_id=1,
        trace_id="trace-mounted-1",
        topic_workspace_id=topic.id,
        source_system="feishu",
        source_message_id="om_123",
        mode_before="plan",
        mode_after="clarify",
        active_agent="Clarifier",
        trigger_reason="missing_owner",
        context_layers_hit={"thread": True, "topic": True},
        final_reply="还差负责人信息。",
    )
    session.add_all([draft, run])
    session.flush()

    trace = MountedAgentTrace(
        tenant_id=1,
        conversation_run_id=run.id,
        agent_name="Clarifier",
        status="completed",
        latency_ms=128,
        input_summary={"mode": "plan"},
        output_summary={"question": "这件事打算谁来负责？"},
    )
    policy = MountedToolCapabilityPolicy(
        tenant_id=1,
        tool_name="project.create",
        risk_level="high",
        source_scope={"systems": ["feishu"]},
    )
    activation = MountedSourceActivation(
        tenant_id=1,
        source_system="feishu",
        scope_type="chat",
        scope_ref="oc_runtime_lab",
        status="active",
        cutover_mode="mounted_only",
        auto_reply_enabled=True,
    )
    session.add_all([trace, activation, policy])
    session.commit()

    saved_topic = session.query(MountedTopicWorkspace).one()
    saved_draft = session.query(MountedDraftWorkspace).one()
    saved_run = session.query(MountedConversationRun).one()
    saved_trace = session.query(MountedAgentTrace).one()
    saved_activation = session.query(MountedSourceActivation).one()
    saved_policy = session.query(MountedToolCapabilityPolicy).one()

    assert saved_topic.status == "active"
    assert saved_draft.state == "collecting"
    assert saved_draft.active_work_item_id == 501
    assert saved_draft.latest_source_message_id == "om_reply_source_1"
    assert saved_draft.latest_reply_message_id == "om_reply_sent_1"
    assert saved_run.status == "queued"
    assert saved_trace.tool_trace is None
    assert saved_activation.cutover_mode == "mounted_only"
    assert saved_activation.auto_reply_enabled is True
    assert saved_policy.confirm_required is True
    assert saved_policy.auto_execute_allowed is False
    assert saved_policy.is_active is True


def test_mounted_ai_runtime_models_enforce_unique_scopes() -> None:
    session = _create_session()

    session.add(
        MountedTopicWorkspace(
            tenant_id=9,
            topic_key="feishu:chat_9:topic_9",
            source_system="feishu",
        )
    )
    session.add(
        MountedSourceActivation(
            tenant_id=9,
            source_system="feishu",
            scope_type="chat",
            scope_ref="oc_runtime_lab",
        )
    )
    session.add(
        MountedToolCapabilityPolicy(
            tenant_id=9,
            tool_name="project.create",
            risk_level="high",
        )
    )
    session.commit()

    session.add(
        MountedTopicWorkspace(
            tenant_id=9,
            topic_key="feishu:chat_9:topic_9",
            source_system="feishu",
        )
    )
    with pytest.raises(IntegrityError):
        session.commit()
    session.rollback()

    session.add(
        MountedSourceActivation(
            tenant_id=9,
            source_system="feishu",
            scope_type="chat",
            scope_ref="oc_runtime_lab",
        )
    )
    with pytest.raises(IntegrityError):
        session.commit()
    session.rollback()

    session.add(
        MountedToolCapabilityPolicy(
            tenant_id=9,
            tool_name="project.create",
            risk_level="high",
        )
    )
    with pytest.raises(IntegrityError):
        session.commit()