test_smoke.py

"""Smoke test：上線前最後檢查（教材第 16.9 節）。

執行：
    pytest tests/test_smoke.py -v

或不裝 pytest 直接跑：
    python -m tests.test_smoke

前置條件：
    1. 已設定 OPENAI_API_KEY
    2. 已執行 python -m src.ingest
"""

from __future__ import annotations

import sys

import pytest

from src.config import config
from src.rag import build_chain


def _ready() -> bool:
    """檢查環境是否準備好可以跑 smoke test。"""
    if not config.has_openai_key:
        return False
    if not config.db_dir.exists() or not any(config.db_dir.iterdir()):
        return False
    return True


pytestmark = pytest.mark.skipif(
    not _ready(),
    reason="需要 OPENAI_API_KEY 與已索引的向量資料庫（先跑 python -m src.ingest）",
)


@pytest.fixture(scope="module")
def rag_chain():
    chain, _ = build_chain()
    return chain


# (問題, 期待答案中包含的關鍵字之一)
SMOKE_CASES = [
    ("公司年假有幾天？", ["14"]),
    ("可以在家工作嗎？", ["在家", "WFH", "遠端"]),
    ("薪資什麼時候發放？", ["5", "五"]),
    ("DataPilot Pro 有哪些方案？", ["Starter", "Business", "Enterprise"]),
    # 不存在的問題：應該回答「資料不足」之類，**不該編造答案**
    ("公司提供股票選擇權嗎？", ["無法回答", "資料不足", "沒有", "未提及"]),
]


@pytest.mark.parametrize("question,keywords", SMOKE_CASES)
def test_smoke(rag_chain, question: str, keywords: list[str]) -> None:
    answer = rag_chain.invoke(question)
    assert answer, f"答案不可為空：{question}"

    matched = any(kw in answer for kw in keywords)
    assert matched, (
        f"\nQ: {question}\nA: {answer}\n"
        f"Expected one of: {keywords}"
    )


def main_cli() -> None:
    """讓使用者可以直接 python -m tests.test_smoke 跑（不裝 pytest 也行）。"""
    if not _ready():
        print("[錯誤] 環境未就緒：請設定 OPENAI_API_KEY 並執行 python -m src.ingest")
        sys.exit(1)

    chain, _ = build_chain()
    failed = 0
    for q, keywords in SMOKE_CASES:
        ans = chain.invoke(q)
        ok = any(kw in ans for kw in keywords)
        status = "PASS" if ok else "FAIL"
        print(f"[{status}] {q}")
        print(f"        -> {ans[:120]}{'...' if len(ans) > 120 else ''}\n")
        if not ok:
            failed += 1

    print(f"\n=== 完成：{len(SMOKE_CASES) - failed}/{len(SMOKE_CASES)} 通過 ===")
    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    main_cli()