влад
@@ -0,0 +1,182 @@
|
||||
"""Юнит-тесты 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()
|
||||
Reference in New Issue
Block a user