"""Юнит-тесты 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()