atlasbot/tests/test_engine.py

"""Answer-engine regression tests."""

from __future__ import annotations

import asyncio
from dataclasses import replace

from atlasbot.engine.answerer import AnswerEngine
from atlasbot.knowledge.loader import KnowledgeBase
from atlasbot.snapshot.builder import SnapshotProvider
from testing.fakes import FakeLLM, SlowFakeLLM, build_test_settings


def test_engine_answer_basic() -> None:
    """The quick path should answer from the fact sheet."""

    llm = FakeLLM()
    settings = build_test_settings()
    kb = KnowledgeBase("")
    snapshot = SnapshotProvider(settings)
    engine = AnswerEngine(settings, llm, kb, snapshot)

    result = asyncio.run(engine.answer("What is Atlas?", mode="quick"))
    assert "Atlas has 22 nodes" in result.reply


def test_smart_mode_uses_factsheet_path() -> None:
    """Smart mode should stay on the factsheet branch for direct cluster questions."""

    llm = FakeLLM()
    settings = build_test_settings()
    kb = KnowledgeBase("")
    snapshot = SnapshotProvider(settings)
    engine = AnswerEngine(settings, llm, kb, snapshot)

    result = asyncio.run(engine.answer("What is the most demanding system in titan lab currently?", mode="smart"))
    assert "Atlas has 22 nodes" in result.reply
    assert "time budget" not in result.reply.lower()


def test_genius_mode_uses_factsheet_path() -> None:
    """Genius mode should also return the factsheet answer for the same query."""

    llm = FakeLLM()
    settings = build_test_settings()
    kb = KnowledgeBase("")
    snapshot = SnapshotProvider(settings)
    engine = AnswerEngine(settings, llm, kb, snapshot)

    result = asyncio.run(engine.answer("What is the most demanding system in titan lab currently?", mode="genius"))
    assert "Atlas has 22 nodes" in result.reply
    assert "time budget" not in result.reply.lower()


def test_plain_math_question_is_rejected_for_cluster_modes() -> None:
    """The bot should keep users on cluster questions instead of generic math."""

    llm = FakeLLM()
    settings = build_test_settings()
    kb = KnowledgeBase("")
    snapshot = SnapshotProvider(settings)
    engine = AnswerEngine(settings, llm, kb, snapshot)

    result = asyncio.run(engine.answer("what is 2+2?", mode="quick"))
    assert "focus on Titan cluster operations" in result.reply


def test_quick_mode_time_budget_guard() -> None:
    """A slow model call should trip the quick-mode budget guard."""

    llm = SlowFakeLLM()
    settings = replace(build_test_settings(), quick_time_budget_sec=0.01)
    kb = KnowledgeBase("")
    snapshot = SnapshotProvider(settings)
    engine = AnswerEngine(settings, llm, kb, snapshot)

    result = asyncio.run(engine.answer("What is Atlas?", mode="quick"))
    assert "time budget" in result.reply
    assert result.meta.get("time_budget_hit") is True