site/WESP_REL/tests/test_admin_llm_tools.py

"""Юнит-тесты slim-tools (без вызова LLM)."""

from __future__ import annotations

import unittest

from app.services.admin_llm_tool_router import _validate_select_only
from app.services.admin_llm_orchestrator import (
    _auto_probe_plan,
    _db_glimpse_table_list_only,
    _extract_dispenser_hint,
    _is_components_question,
    _looks_unreliable_final,
    _structured_answer_from_trace,
)


class AdminLlmSqlPolicyTests(unittest.TestCase):
    def test_accepts_simple_select(self) -> None:
        s = _validate_select_only("  SELECT 1 AS x  ")
        self.assertIn("SELECT", s.upper())

    def test_accepts_with_cte(self) -> None:
        s = _validate_select_only("WITH a AS (SELECT 1) SELECT * FROM a")
        self.assertTrue(s.upper().startswith("WITH"))

    def test_rejects_semicolon_two_statements(self) -> None:
        with self.assertRaises(ValueError):
            _validate_select_only("SELECT 1; SELECT 2")

    def test_rejects_insert(self) -> None:
        with self.assertRaises(ValueError):
            _validate_select_only("INSERT INTO t VALUES (1)")

    def test_rejects_pragma(self) -> None:
        with self.assertRaises(ValueError):
            _validate_select_only("PRAGMA integrity_check")


class AdminLlmOrchestratorHelpersTests(unittest.TestCase):
    def test_extract_dispenser_hint_after_marker(self) -> None:
        self.assertEqual(
            _extract_dispenser_hint("какие есть для кормораздатчика лох"),
            "лох",
        )

    def test_extract_dispenser_hint_empty_for_list_in_db(self) -> None:
        self.assertEqual(_extract_dispenser_hint("какие кормораздатчики есть в бд?"), "")

    def test_extract_dispenser_hint_quoted(self) -> None:
        self.assertEqual(_extract_dispenser_hint('рецепты для «шкаф А»'), "шкаф А")

    def test_auto_probe_includes_distinct_recipes_when_named_feeder_without_word_recipe(self) -> None:
        plan = _auto_probe_plan("какие есть для кормораздатчика лох")
        sqls = [a[1].get("sql", "") for a in plan if a[0] == "slim_sql"]
        self.assertTrue(any("DISTINCT r.name AS recipe" in s and "лох" in s for s in sqls))

    def test_components_question_typo_still_detected(self) -> None:
        self.assertTrue(_is_components_question("кикие есть компонентв".lower(), "кикие есть компонентв"))

    def test_auto_probe_db_glimpse_runs_sqlite_master(self) -> None:
        plan = _auto_probe_plan("ок, а что есть бд?")
        sqls = [a[1].get("sql", "") for a in plan if a[0] == "slim_sql"]
        self.assertTrue(any("sqlite_master" in s for s in sqls))

    def test_db_glimpse_table_list_skips_feeder_questions(self) -> None:
        self.assertTrue(_db_glimpse_table_list_only("что есть в бд"))
        self.assertFalse(_db_glimpse_table_list_only("какие кормораздатчики есть в бд"))

    def test_looks_unreliable_detects_hallucinated_db_files(self) -> None:
        self.assertTrue(
            _looks_unreliable_final(
                "что в базе",
                "В recipes.db хранятся рецепты, в reports.db — отчёты.",
            )
        )

    def test_structured_feeders_not_replaced_by_table_catalog(self) -> None:
        trace = [
            {
                "tool": "slim_sql",
                "ok": True,
                "result_preview": {
                    "rows_sample": [{"name": "recipe"}, {"name": "feed_dispenser"}],
                    "row_keys": ["name"],
                },
            },
            {
                "tool": "slim_sql",
                "ok": True,
                "result_preview": {
                    "rows_sample": [
                        {"id": "1", "name": "Шкаф 1", "farm": "Ферма А", "is_active": 1},
                    ],
                    "row_keys": ["farm", "id", "is_active", "name"],
                },
            },
        ]
        out = _structured_answer_from_trace("какие кормораздатчики есть в бд?", trace)
        self.assertIn("Шкаф 1", out)
        self.assertTrue(out.startswith("Кормораздатчики"))

    def test_auto_probe_includes_component_selects(self) -> None:
        plan = _auto_probe_plan("какие компоненты в бд")
        sqls = [a[1].get("sql", "") for a in plan if a[0] == "slim_sql"]
        self.assertTrue(any("pragma_table_info('component')" in s for s in sqls))
        self.assertTrue(any("FROM component" in s and "dry_matter" in s for s in sqls))
        self.assertTrue(any("sqlite_master" in s for s in sqls))

    def test_structured_answer_prefers_component_rows_not_sqlite_master(self) -> None:
        trace = [
            {
                "tool": "slim_sql",
                "ok": True,
                "result_preview": {
                    "rows_sample": [{"name": "recipe"}, {"name": "component"}],
                    "row_keys": ["name"],
                },
            },
            {
                "tool": "slim_sql",
                "ok": True,
                "result_preview": {
                    "rows_sample": [
                        {
                            "name": "Соя",
                            "type": "кг",
                            "is_active": 1,
                            "is_deleted": 0,
                            "dry_matter": 1.0,
                            "protein": 2.0,
                            "energy": 3.0,
                            "price": 0.0,
                        }
                    ],
                    "row_keys": [
                        "dry_matter",
                        "energy",
                        "is_active",
                        "is_deleted",
                        "name",
                        "price",
                        "protein",
                        "type",
                    ],
                },
            },
        ]
        out = _structured_answer_from_trace("что в таблице component", trace)
        self.assertIn("Соя", out)
        self.assertNotIn("alembic", out)

    def test_structured_answer_uses_recipe_column_not_dispenser_names(self) -> None:
        trace = [
            {
                "tool": "slim_sql",
                "ok": True,
                "result_preview": {
                    "rows_sample": [
                        {"dispenser": "лох", "farm": "X", "operator": "Иванов", "period": "утро", "recipe": "Рацион А"},
                        {"dispenser": "другой", "farm": "Y", "operator": "", "period": "утро", "recipe": "Рацион Б"},
                    ],
                    "row_keys": ["dispenser", "farm", "operator", "period", "recipe"],
                },
            },
            {
                "tool": "slim_sql",
                "ok": True,
                "result_preview": {
                    "rows_sample": [{"id": "1", "name": "лох", "farm": "X", "is_active": 1}],
                    "row_keys": ["farm", "id", "is_active", "name"],
                },
            },
        ]
        out = _structured_answer_from_trace("рецепты для кормораздатчика лох", trace)
        self.assertIn("Рацион А", out)
        self.assertNotIn("Иванов", out)
        self.assertNotIn("Рацион Б", out)


if __name__ == "__main__":
    unittest.main()