aiagent/backend/tests/test_company_orchestrator.py

"""
虚拟公司编排器测试 — 全程 mock LLM，不产生真实 API 调用

覆盖：
- _parse_ceo_json：CEO 规划输出的 JSON 提取（直接/代码块/尾逗号/垃圾输入）
- _select_relevant_knowledge：学习闭环的知识相关性筛选（纯函数）
- _build_dept_context：四档-① 部门间交接的上下文构造与字数上限
- execute_stream：波次 DAG 执行顺序、部门间交接注入、per-dept 独立 session
- _ceo_review_with_scoring：评审 JSON 解析失败的保守兜底
"""
import json
from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest

# 注册公司相关模型到 Base.metadata（db_session fixture 建表用）
import app.models.company  # noqa: F401
import app.models.company_knowledge  # noqa: F401
import app.models.team  # noqa: F401
from app.models.agent import Agent
from app.models.company import Company
from app.models.team import Team

from app.services import company_orchestrator as co_module
from app.services.company_orchestrator import (
    CompanyOrchestrator,
    _parse_ceo_json,
    _select_relevant_knowledge,
)


# ─── _parse_ceo_json ───

class TestParseCeoJson:
    def test_direct_json(self):
        plan = {"analysis": "a", "departments": [{"department_name": "数据组"}]}
        assert _parse_ceo_json(json.dumps(plan, ensure_ascii=False)) == plan

    def test_markdown_fenced(self):
        plan = {"analysis": "a", "departments": []}
        raw = f"```json\n{json.dumps(plan, ensure_ascii=False)}\n```"
        assert _parse_ceo_json(raw) == plan

    def test_trailing_comma_fixed(self):
        raw = '{"analysis": "a", "departments": [{"department_name": "数据组",}],}'
        result = _parse_ceo_json(raw)
        assert result["departments"][0]["department_name"] == "数据组"

    def test_garbage_returns_empty(self):
        assert _parse_ceo_json("完全不是 JSON 的输出") == {}

    def test_brace_extract_with_surrounding_text(self):
        raw = '这是计划：\n{"analysis": "a", "departments": []}\n以上。'
        result = _parse_ceo_json(raw)
        assert result.get("analysis") == "a"


# ─── _select_relevant_knowledge ───

class TestSelectRelevantKnowledge:
    def test_empty_entries(self):
        assert _select_relevant_knowledge([], "市场营销", 5, 1000) == ""

    def test_zero_limits(self):
        entries = [{"title": "市场", "content": "市场", "category": "review", "tags": []}]
        assert _select_relevant_knowledge(entries, "市场", 0, 1000) == ""
        assert _select_relevant_knowledge(entries, "市场", 5, 0) == ""

    def test_irrelevant_excluded(self):
        entries = [
            {"title": "市场分析教训", "content": "上次市场项目的增长策略", "category": "review",
             "tags": ["市场"], "_recency": 2},
            {"title": "无关条目", "content": "烹饪菜谱大全", "category": "general",
             "tags": [], "_recency": 1},
        ]
        result = _select_relevant_knowledge(entries, "做一个市场营销产品", 5, 2000)
        assert "市场分析教训" in result
        assert "无关条目" not in result

    def test_no_relevant_returns_empty(self):
        entries = [{"title": "菜谱", "content": "红烧肉做法", "category": "general", "tags": []}]
        assert _select_relevant_knowledge(entries, "做一个市场营销产品", 5, 2000) == ""

    def test_max_chars_truncation(self):
        entries = [
            {"title": "市场教训", "content": "市场" + "长" * 500, "category": "review", "tags": []},
        ]
        result = _select_relevant_knowledge(entries, "市场分析", 5, 50)
        assert result != ""
        # 正文块被截到 50 字以内（外加固定头部）
        assert "长" * 100 not in result

    def test_review_category_weighted_higher(self):
        """review/insight 类别权重高，同等主题命中下排前。"""
        entries = [
            {"title": "部门产出", "content": "市场 数据", "category": "dept_output",
             "tags": [], "_recency": 99},
            {"title": "审查教训", "content": "市场 数据", "category": "review",
             "tags": [], "_recency": 1},
        ]
        result = _select_relevant_knowledge(entries, "市场数据分析", 1, 2000)
        # 只取 1 条时，权重更高的 review 入选
        assert "审查教训" in result
        assert "部门产出" not in result


# ─── _build_dept_context ───

class TestBuildDeptContext:
    def _orch(self):
        return CompanyOrchestrator(db=None, company_id="c1", user_id="u1")

    def test_no_dependencies_returns_base(self):
        orch = self._orch()
        ctx = orch._build_dept_context({"dependencies": []}, "CEO 分析", {"A": "交付物"})
        assert ctx == "CEO 分析"

    def test_handoff_injected(self):
        orch = self._orch()
        ctx = orch._build_dept_context(
            {"dependencies": ["数据组"]}, "CEO 分析", {"数据组": "上游数据内容"}
        )
        assert "上游部门「数据组」" in ctx
        assert "上游数据内容" in ctx
        assert "CEO 分析" in ctx

    def test_per_dept_char_limit(self):
        orch = self._orch()
        ctx = orch._build_dept_context(
            {"dependencies": ["A"]}, "", {"A": "x" * 5000}
        )
        # 单部门交付物注入上限 2000 字
        assert "x" * 2000 in ctx
        assert "x" * 2001 not in ctx

    def test_total_char_limit(self):
        orch = self._orch()
        deliverables = {n: "x" * 2000 for n in ["A", "B", "C", "D"]}
        ctx = orch._build_dept_context(
            {"dependencies": ["A", "B", "C", "D"]}, "", deliverables
        )
        # 合计上限 6000：前三个各 2000 注满，第四个被截断
        assert "上游部门「A」" in ctx
        assert "上游部门「C」" in ctx
        assert "上游部门「D」" not in ctx

    def test_missing_deliverable_skipped(self):
        orch = self._orch()
        ctx = orch._build_dept_context(
            {"dependencies": ["不存在"]}, "CEO 分析", {}
        )
        assert ctx == "CEO 分析"


# ─── execute_stream 全流程（mock TeamOrchestrator + CEO 评审）───

class _FakeTeamOrchestrator:
    """替代 TeamOrchestrator：记录收到的 db/task，吐固定事件流，不调 LLM。"""

    instances = []

    def __init__(self, db, team_id, user_id, auto_approve_files=True):
        self.db = db
        self.team_id = team_id
        self.tasks = []
        _FakeTeamOrchestrator.instances.append(self)

    async def execute_stream(self, task, auto_approve_files=True):
        self.tasks.append(task)
        yield {"type": "phase_start", "phase": 1, "total_phases": 1, "name": "执行"}
        yield {
            "type": "final",
            "deliverable": f"交付物-{self.team_id}",
            "files": [f"{self.team_id}.py"],
            "phase_count": 1,
        }


def _make_company(db, dept_names):
    company = Company(name="测试公司", user_id="u1")
    db.add(company)
    db.flush()
    teams = []
    for i, name in enumerate(dept_names):
        t = Team(name=name, user_id="u1", parent_company_id=company.id,
                 department_type=f"type{i}")
        db.add(t)
        teams.append(t)
    db.commit()
    return company, teams


def _patch_session_local(monkeypatch):
    """把编排器里的 SessionLocal 换成 MagicMock 工厂，追踪 per-dept session 生命周期。"""
    sessions = []

    def factory():
        s = MagicMock(name=f"dept_session_{len(sessions)}")
        sessions.append(s)
        return s

    monkeypatch.setattr(co_module, "SessionLocal", factory)
    return sessions


def _patch_team_orchestrator(monkeypatch):
    import app.services.team_orchestrator as to_module
    _FakeTeamOrchestrator.instances = []
    monkeypatch.setattr(to_module, "TeamOrchestrator", _FakeTeamOrchestrator)
    return _FakeTeamOrchestrator.instances


def _patch_review(monkeypatch, all_pass=True):
    async def fake(self, company, ceo_plan, dept_results, project_desc):
        return {
            "summary": "整体良好",
            "overall_score": 8,
            "next_steps": "",
            "departments": [
                {"name": r["department_name"], "score": 8 if all_pass else 4,
                 "feedback": "fb", "pass": all_pass}
                for r in dept_results
            ],
        }
    monkeypatch.setattr(CompanyOrchestrator, "_ceo_review_with_scoring", fake)


def _disable_knowledge_extract(monkeypatch):
    from app.core.config import settings
    monkeypatch.setattr(settings, "COMPANY_AUTO_EXTRACT_KNOWLEDGE", False)


async def test_execute_stream_parallel_wave_per_dept_sessions(db_session, monkeypatch, tmp_path):
    """同波并行部门：各自拿到独立 session（非编排器共享 session），且用完关闭。"""
    monkeypatch.chdir(tmp_path)  # project_path mkdir 落到临时目录
    company, teams = _make_company(db_session, ["数据组", "前端组"])
    plan = {"analysis": "分析", "departments": [
        {"department_name": "数据组", "goal": "g1", "deliverables": ["a"], "dependencies": []},
        {"department_name": "前端组", "goal": "g2", "deliverables": ["b"], "dependencies": []},
    ]}
    sessions = _patch_session_local(monkeypatch)
    instances = _patch_team_orchestrator(monkeypatch)
    _patch_review(monkeypatch)
    _disable_knowledge_extract(monkeypatch)

    orch = CompanyOrchestrator(db_session, company.id, "u1")
    events = [e async for e in orch.execute_stream("测试项目", ceo_plan=plan)]

    types = [e["type"] for e in events]
    assert "company_done" in types
    assert types.count("dept_done") >= 2

    # 核心断言：两个部门各自新建 session，且都已关闭
    assert len(sessions) == 2
    assert sessions[0] is not sessions[1]
    assert all(s.close.called for s in sessions)
    # TeamOrchestrator 收到的不是编排器共享的 self.db
    assert len(instances) == 2
    assert all(inst.db is not db_session for inst in instances)
    assert all(inst.db in sessions for inst in instances)

    done = [e for e in events if e["type"] == "company_done"][0]
    assert done["success"] is True


async def test_execute_stream_dependency_order_and_handoff(db_session, monkeypatch, tmp_path):
    """依赖 DAG：下游等上游完成后才启动，且上游交付物注入下游任务上下文。"""
    monkeypatch.chdir(tmp_path)
    company, teams = _make_company(db_session, ["数据组", "前端组"])
    plan = {"analysis": "分析", "departments": [
        {"department_name": "数据组", "goal": "g1", "deliverables": ["a"], "dependencies": []},
        {"department_name": "前端组", "goal": "g2", "deliverables": ["b"],
         "dependencies": ["数据组"]},
    ]}
    _patch_session_local(monkeypatch)
    instances = _patch_team_orchestrator(monkeypatch)
    _patch_review(monkeypatch)
    _disable_knowledge_extract(monkeypatch)

    orch = CompanyOrchestrator(db_session, company.id, "u1")
    events = [e async for e in orch.execute_stream("测试项目", ceo_plan=plan)]

    # 第一波：前端组依赖未满足 → dept_waiting
    assert any(e["type"] == "dept_waiting" and e.get("department_name") == "前端组"
               for e in events)

    # 数据组先实例化执行，前端组第二波才启动
    assert len(instances) == 2
    assert instances[0].team_id == teams[0].id
    assert instances[1].team_id == teams[1].id

    # 四档-① 交接：数据组的交付物注入前端组任务上下文
    downstream_task = instances[1].tasks[0]
    assert "上游部门「数据组」" in downstream_task
    assert f"交付物-{teams[0].id}" in downstream_task


async def test_execute_stream_uses_shared_session_only_sequentially(
        db_session, monkeypatch, tmp_path):
    """编排主流程（建项目/落状态）仍用 self.db，但绝不把它交给并行部门。"""
    monkeypatch.chdir(tmp_path)
    company, teams = _make_company(db_session, ["数据组"])
    plan = {"analysis": "分析", "departments": [
        {"department_name": "数据组", "goal": "g", "deliverables": ["a"], "dependencies": []},
    ]}
    sessions = _patch_session_local(monkeypatch)
    instances = _patch_team_orchestrator(monkeypatch)
    _patch_review(monkeypatch)
    _disable_knowledge_extract(monkeypatch)

    orch = CompanyOrchestrator(db_session, company.id, "u1")
    events = [e async for e in orch.execute_stream("测试项目", ceo_plan=plan)]

    assert any(e["type"] == "company_done" for e in events)
    assert len(instances) == 1
    assert instances[0].db is sessions[0]
    assert instances[0].db is not db_session


# ─── _ceo_review_with_scoring 解析兜底 ───

class _FakeLLMClient:
    """替代 _LLMClient：返回固定 content，不调真实 API。"""
    content = ""

    def __init__(self, cfg):
        pass

    async def chat(self, messages, tools, iteration):
        return SimpleNamespace(content=_FakeLLMClient.content)


def _setup_ceo_agent(db_session):
    agent = Agent(name="测试CEO", workflow_config={}, user_id="u1")
    db_session.add(agent)
    db_session.flush()
    company = Company(name="测试公司", user_id="u1", ceo_agent_id=agent.id)
    db_session.add(company)
    db_session.commit()
    return company


async def test_ceo_review_valid_json(db_session, monkeypatch):
    company = _setup_ceo_agent(db_session)
    payload = {
        "overall_score": 8, "summary": "好", "next_steps": "继续",
        "departments": [{"name": "数据组", "score": 8, "feedback": "好",
                         "pass": True, "improvement_areas": []}],
    }
    _FakeLLMClient.content = json.dumps(payload, ensure_ascii=False)
    import app.agent_runtime.core as arc
    monkeypatch.setattr(arc, "_LLMClient", _FakeLLMClient)

    orch = CompanyOrchestrator(db_session, company.id, "u1")
    review = await orch._ceo_review_with_scoring(
        company, {"departments": []},
        [{"department_name": "数据组", "success": True, "result": {"deliverable": "x"}}],
        "项目",
    )
    assert review["overall_score"] == 8
    assert review["departments"][0]["pass"] is True
    assert not review.get("parse_failed")


async def test_ceo_review_parse_failure_conservative(db_session, monkeypatch):
    """评审输出两次都不是合法 JSON → 保守判未通过 + 显式 parse_failed，绝不静默 auto-pass。"""
    company = _setup_ceo_agent(db_session)
    _FakeLLMClient.content = "这根本不是 JSON，两次重试也都不是"
    import app.agent_runtime.core as arc
    monkeypatch.setattr(arc, "_LLMClient", _FakeLLMClient)

    orch = CompanyOrchestrator(db_session, company.id, "u1")
    review = await orch._ceo_review_with_scoring(
        company, {"departments": []},
        [{"department_name": "数据组", "success": True, "result": {"deliverable": "x"}}],
        "项目",
    )
    assert review.get("parse_failed") is True
    assert review["departments"][0]["pass"] is False


async def test_ceo_review_llm_exception_conservative(db_session, monkeypatch):
    """LLM 调用直接抛异常 → 同样走保守兜底。"""
    company = _setup_ceo_agent(db_session)

    class _BoomClient:
        def __init__(self, cfg):
            pass

        async def chat(self, messages, tools, iteration):
            raise RuntimeError("network down")

    import app.agent_runtime.core as arc
    monkeypatch.setattr(arc, "_LLMClient", _BoomClient)

    orch = CompanyOrchestrator(db_session, company.id, "u1")
    review = await orch._ceo_review_with_scoring(
        company, {"departments": []},
        [{"department_name": "数据组", "success": True, "result": {"deliverable": "x"}}],
        "项目",
    )
    assert review.get("parse_failed") is True
    assert review["departments"][0]["pass"] is False