530 lines
22 KiB
Python
530 lines
22 KiB
Python
import json
|
||
import os
|
||
import tempfile
|
||
import unittest
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
from werkzeug.security import generate_password_hash
|
||
|
||
from app import create_app, db
|
||
from app.models import WebUser
|
||
from app.services.admin_llm_activity_log import (
|
||
append_admin_llm_activity,
|
||
format_llm_log_line_text,
|
||
read_admin_llm_activity,
|
||
)
|
||
from app.services.admin_llm_rate_limit import reset_llm_rate_limit_for_tests
|
||
from tests.test_admin_panel_access_and_users import AdminPanelTestConfig
|
||
|
||
|
||
class AdminLLMEnabledConfig(AdminPanelTestConfig):
|
||
WESP_ADMIN_LLM_ENABLED = True
|
||
|
||
|
||
class AdminLLMToolsEnabledConfig(AdminLLMEnabledConfig):
|
||
WESP_LLM_TOOLS_ENABLED = True
|
||
|
||
|
||
def _openai_chat_response(text: str) -> dict:
|
||
return {
|
||
"choices": [
|
||
{
|
||
"message": {"role": "assistant", "content": text},
|
||
}
|
||
]
|
||
}
|
||
|
||
|
||
class AdminLLMEndpointsTests(unittest.TestCase):
|
||
def setUp(self) -> None:
|
||
reset_llm_rate_limit_for_tests()
|
||
self._prev_security_settings_path = os.environ.get("WESP_SECURITY_SETTINGS_PATH")
|
||
self._prev_network_diagnostics_path = os.environ.get("WESP_NETWORK_DIAGNOSTICS_PATH")
|
||
_sec_dir = tempfile.mkdtemp(prefix="wesp-llm-sec-")
|
||
os.environ["WESP_SECURITY_SETTINGS_PATH"] = os.path.join(_sec_dir, "wesp_security_settings.json")
|
||
_diag_dir = tempfile.mkdtemp(prefix="wesp-llm-diag-")
|
||
os.environ["WESP_NETWORK_DIAGNOSTICS_PATH"] = os.path.join(
|
||
_diag_dir, "wesp_network_diagnostics.json"
|
||
)
|
||
|
||
self.app = create_app(AdminLLMEnabledConfig)
|
||
self.client = self.app.test_client()
|
||
self.ctx = self.app.app_context()
|
||
self.ctx.push()
|
||
os.makedirs(AdminLLMEnabledConfig.DATA_DIR, exist_ok=True)
|
||
db.create_all()
|
||
db.session.add(
|
||
WebUser(
|
||
login="admin",
|
||
password_hash=generate_password_hash("admin-secret"),
|
||
is_superuser=True,
|
||
)
|
||
)
|
||
db.session.add(
|
||
WebUser(
|
||
login="operator",
|
||
password_hash=generate_password_hash("operator-secret"),
|
||
is_superuser=False,
|
||
)
|
||
)
|
||
db.session.commit()
|
||
|
||
def tearDown(self) -> None:
|
||
db.session.remove()
|
||
db.drop_all()
|
||
self.ctx.pop()
|
||
reset_llm_rate_limit_for_tests()
|
||
if self._prev_security_settings_path is None:
|
||
os.environ.pop("WESP_SECURITY_SETTINGS_PATH", None)
|
||
else:
|
||
os.environ["WESP_SECURITY_SETTINGS_PATH"] = self._prev_security_settings_path
|
||
if self._prev_network_diagnostics_path is None:
|
||
os.environ.pop("WESP_NETWORK_DIAGNOSTICS_PATH", None)
|
||
else:
|
||
os.environ["WESP_NETWORK_DIAGNOSTICS_PATH"] = self._prev_network_diagnostics_path
|
||
|
||
def _login(self, login: str, password: str) -> None:
|
||
r = self.client.post("/api/auth/login", json={"login": login, "password": password})
|
||
self.assertEqual(r.status_code, 200)
|
||
|
||
def test_operator_forbidden_llm_status(self) -> None:
|
||
self._login("operator", "operator-secret")
|
||
r = self.client.get("/api/admin/llm/status")
|
||
self.assertEqual(r.status_code, 403)
|
||
|
||
def test_superuser_llm_status_when_disabled_config(self) -> None:
|
||
tmp = tempfile.mkdtemp(prefix="wesp-llm-disabled-")
|
||
|
||
class _Isolated(AdminPanelTestConfig):
|
||
BASE_DIR = tmp
|
||
DATA_DIR = os.path.join(tmp, "data")
|
||
SQLALCHEMY_DATABASE_URI = f"sqlite:///{os.path.join(tmp, 'recipes_llm_dis.db')}"
|
||
SQLALCHEMY_BINDS = {"reports": f"sqlite:///{os.path.join(tmp, 'reports_llm_dis.db')}"}
|
||
|
||
app = create_app(_Isolated)
|
||
try:
|
||
os.makedirs(_Isolated.DATA_DIR, exist_ok=True)
|
||
with app.test_client() as c:
|
||
with app.app_context():
|
||
db.create_all()
|
||
db.session.add(
|
||
WebUser(
|
||
login="admin",
|
||
password_hash=generate_password_hash("admin-secret"),
|
||
is_superuser=True,
|
||
)
|
||
)
|
||
db.session.commit()
|
||
c.post("/api/auth/login", json={"login": "admin", "password": "admin-secret"})
|
||
r = c.get("/api/admin/llm/status")
|
||
self.assertEqual(r.status_code, 200)
|
||
data = r.get_json()
|
||
self.assertFalse(data.get("enabled"))
|
||
finally:
|
||
with app.app_context():
|
||
db.session.remove()
|
||
db.drop_all()
|
||
|
||
@patch("app.services.local_llm_client.requests.get")
|
||
def test_llm_status_models_ok(self, mock_get) -> None:
|
||
mock_get.return_value.status_code = 200
|
||
mock_get.return_value.json.return_value = {
|
||
"data": [{"id": "qwen2.5-1.5b-instruct-q5_k_m.gguf"}]
|
||
}
|
||
self._login("admin", "admin-secret")
|
||
r = self.client.get("/api/admin/llm/status")
|
||
self.assertEqual(r.status_code, 200)
|
||
data = r.get_json()
|
||
self.assertTrue(data.get("enabled"))
|
||
self.assertTrue(data.get("llm_reachable"))
|
||
self.assertTrue(data.get("model_present"))
|
||
|
||
@patch("app.services.local_llm_client.requests.post")
|
||
def test_llm_summary_mock_chat(self, mock_post) -> None:
|
||
mock_post.return_value.status_code = 200
|
||
mock_post.return_value.json.return_value = _openai_chat_response("Кратко: отклонений нет.")
|
||
self._login("admin", "admin-secret")
|
||
r = self.client.post(
|
||
"/api/admin/llm/summary",
|
||
data=json.dumps(
|
||
{"date_from": "2026-01-01", "date_to": "2026-01-07", "focus": "both"}
|
||
),
|
||
content_type="application/json",
|
||
)
|
||
self.assertEqual(r.status_code, 200)
|
||
data = r.get_json()
|
||
self.assertEqual(data.get("reply"), "Кратко: отклонений нет.")
|
||
mock_post.assert_called_once()
|
||
|
||
@patch("app.services.local_llm_client.requests.post")
|
||
def test_llm_ping_mock(self, mock_post) -> None:
|
||
mock_post.return_value.status_code = 200
|
||
mock_post.return_value.json.return_value = _openai_chat_response("ок")
|
||
self._login("admin", "admin-secret")
|
||
r = self.client.post("/api/admin/llm/ping", json={})
|
||
self.assertEqual(r.status_code, 200)
|
||
self.assertEqual(r.get_json().get("reply"), "ок")
|
||
mock_post.assert_called_once()
|
||
|
||
@patch("app.services.local_llm_client.requests.post")
|
||
def test_llm_chat_mock(self, mock_post) -> None:
|
||
mock_post.return_value.status_code = 200
|
||
mock_post.return_value.json.return_value = _openai_chat_response("Здравствуйте.")
|
||
self._login("admin", "admin-secret")
|
||
r = self.client.post(
|
||
"/api/admin/llm/chat",
|
||
data=json.dumps(
|
||
{"messages": [{"role": "user", "content": "Привет"}]}
|
||
),
|
||
content_type="application/json",
|
||
)
|
||
self.assertEqual(r.status_code, 200)
|
||
self.assertEqual(r.get_json().get("reply"), "Здравствуйте.")
|
||
mock_post.assert_called_once()
|
||
body = (mock_post.call_args.kwargs or {}).get("json") or {}
|
||
self.assertNotIn("tools", body)
|
||
|
||
@patch("app.services.local_llm_client.requests.post")
|
||
def test_llm_chat_sends_tools_when_globally_enabled(self, mock_post) -> None:
|
||
tmp = tempfile.mkdtemp(prefix="wesp-llm-tools-on-")
|
||
|
||
class _Cfg(AdminPanelTestConfig):
|
||
BASE_DIR = tmp
|
||
DATA_DIR = os.path.join(tmp, "data")
|
||
SQLALCHEMY_DATABASE_URI = f"sqlite:///{os.path.join(tmp, 'r_tools.db')}"
|
||
SQLALCHEMY_BINDS = {"reports": f"sqlite:///{os.path.join(tmp, 'rep_tools.db')}"}
|
||
WESP_ADMIN_LLM_ENABLED = True
|
||
WESP_LLM_TOOLS_ENABLED = True
|
||
|
||
app = create_app(_Cfg)
|
||
try:
|
||
os.makedirs(_Cfg.DATA_DIR, exist_ok=True)
|
||
with app.test_client() as c:
|
||
with app.app_context():
|
||
db.create_all()
|
||
db.session.add(
|
||
WebUser(
|
||
login="admin",
|
||
password_hash=generate_password_hash("admin-secret"),
|
||
is_superuser=True,
|
||
)
|
||
)
|
||
db.session.commit()
|
||
c.post("/api/auth/login", json={"login": "admin", "password": "admin-secret"})
|
||
mock_post.return_value.status_code = 200
|
||
mock_post.return_value.json.return_value = _openai_chat_response("Ок.")
|
||
r = c.post(
|
||
"/api/admin/llm/chat",
|
||
data=json.dumps({"messages": [{"role": "user", "content": "Тест"}]}),
|
||
content_type="application/json",
|
||
)
|
||
self.assertEqual(r.status_code, 200)
|
||
body = (mock_post.call_args.kwargs or {}).get("json") or {}
|
||
self.assertIn("tools", body)
|
||
self.assertTrue(body.get("tools"))
|
||
finally:
|
||
with app.app_context():
|
||
db.session.remove()
|
||
db.drop_all()
|
||
|
||
def test_llm_summary_503_when_disabled(self) -> None:
|
||
tmp = tempfile.mkdtemp(prefix="wesp-llm-sum503-")
|
||
|
||
class _Iso(AdminPanelTestConfig):
|
||
BASE_DIR = tmp
|
||
DATA_DIR = os.path.join(tmp, "data")
|
||
SQLALCHEMY_DATABASE_URI = f"sqlite:///{os.path.join(tmp, 'r503.db')}"
|
||
SQLALCHEMY_BINDS = {"reports": f"sqlite:///{os.path.join(tmp, 'rep503.db')}"}
|
||
WESP_ADMIN_LLM_ENABLED = False
|
||
|
||
app = create_app(_Iso)
|
||
try:
|
||
os.makedirs(_Iso.DATA_DIR, exist_ok=True)
|
||
with app.test_client() as c:
|
||
with app.app_context():
|
||
db.create_all()
|
||
db.session.add(
|
||
WebUser(
|
||
login="admin",
|
||
password_hash=generate_password_hash("admin-secret"),
|
||
is_superuser=True,
|
||
)
|
||
)
|
||
db.session.commit()
|
||
c.post("/api/auth/login", json={"login": "admin", "password": "admin-secret"})
|
||
r = c.post(
|
||
"/api/admin/llm/summary",
|
||
data=json.dumps(
|
||
{"date_from": "2026-01-01", "date_to": "2026-01-02", "focus": "both"}
|
||
),
|
||
content_type="application/json",
|
||
)
|
||
self.assertEqual(r.status_code, 503)
|
||
finally:
|
||
with app.app_context():
|
||
db.session.remove()
|
||
db.drop_all()
|
||
|
||
|
||
class AdminLLMToolsChatTests(unittest.TestCase):
|
||
"""Чат с slim-tools: многошаговой ответ LLM (два POST при tool_calls)."""
|
||
|
||
def setUp(self) -> None:
|
||
reset_llm_rate_limit_for_tests()
|
||
self._prev_security_settings_path = os.environ.get("WESP_SECURITY_SETTINGS_PATH")
|
||
self._prev_network_diagnostics_path = os.environ.get("WESP_NETWORK_DIAGNOSTICS_PATH")
|
||
_sec_dir = tempfile.mkdtemp(prefix="wesp-llm-tools-sec-")
|
||
os.environ["WESP_SECURITY_SETTINGS_PATH"] = os.path.join(_sec_dir, "wesp_security_settings.json")
|
||
_diag_dir = tempfile.mkdtemp(prefix="wesp-llm-tools-diag-")
|
||
os.environ["WESP_NETWORK_DIAGNOSTICS_PATH"] = os.path.join(
|
||
_diag_dir, "wesp_network_diagnostics.json"
|
||
)
|
||
|
||
self.app = create_app(AdminLLMToolsEnabledConfig)
|
||
self.client = self.app.test_client()
|
||
self.ctx = self.app.app_context()
|
||
self.ctx.push()
|
||
os.makedirs(AdminLLMToolsEnabledConfig.DATA_DIR, exist_ok=True)
|
||
db.create_all()
|
||
db.session.add(
|
||
WebUser(
|
||
login="admin",
|
||
password_hash=generate_password_hash("admin-secret"),
|
||
is_superuser=True,
|
||
)
|
||
)
|
||
db.session.commit()
|
||
|
||
def tearDown(self) -> None:
|
||
db.session.remove()
|
||
db.drop_all()
|
||
self.ctx.pop()
|
||
reset_llm_rate_limit_for_tests()
|
||
if self._prev_security_settings_path is None:
|
||
os.environ.pop("WESP_SECURITY_SETTINGS_PATH", None)
|
||
else:
|
||
os.environ["WESP_SECURITY_SETTINGS_PATH"] = self._prev_security_settings_path
|
||
if self._prev_network_diagnostics_path is None:
|
||
os.environ.pop("WESP_NETWORK_DIAGNOSTICS_PATH", None)
|
||
else:
|
||
os.environ["WESP_NETWORK_DIAGNOSTICS_PATH"] = self._prev_network_diagnostics_path
|
||
|
||
def _login(self) -> None:
|
||
r = self.client.post("/api/auth/login", json={"login": "admin", "password": "admin-secret"})
|
||
self.assertEqual(r.status_code, 200)
|
||
|
||
@patch("app.services.local_llm_client.requests.post")
|
||
def test_llm_chat_tool_loop_sql_preview(self, mock_post) -> None:
|
||
def _post(url: str, **kwargs: object) -> MagicMock:
|
||
body = kwargs.get("json") or {}
|
||
msgs = body.get("messages") or []
|
||
has_tool = any(
|
||
isinstance(m, dict) and m.get("role") == "tool" for m in msgs
|
||
)
|
||
m = MagicMock()
|
||
m.status_code = 200
|
||
if not has_tool:
|
||
m.json.return_value = {
|
||
"choices": [
|
||
{
|
||
"message": {
|
||
"role": "assistant",
|
||
"content": "",
|
||
"tool_calls": [
|
||
{
|
||
"id": "call_1",
|
||
"type": "function",
|
||
"function": {
|
||
"name": "slim_sql",
|
||
"arguments": json.dumps(
|
||
{
|
||
"bind": "recipes",
|
||
"sql": "SELECT 1 AS one",
|
||
"mode": "preview",
|
||
}
|
||
),
|
||
},
|
||
}
|
||
],
|
||
}
|
||
}
|
||
]
|
||
}
|
||
else:
|
||
m.json.return_value = {
|
||
"choices": [
|
||
{"message": {"role": "assistant", "content": "План запроса получен."}}
|
||
]
|
||
}
|
||
return m
|
||
|
||
mock_post.side_effect = _post
|
||
self._login()
|
||
r = self.client.post(
|
||
"/api/admin/llm/chat",
|
||
data=json.dumps({"messages": [{"role": "user", "content": "Покажи план для SELECT 1"}]}),
|
||
content_type="application/json",
|
||
)
|
||
self.assertEqual(r.status_code, 200, r.get_data(as_text=True))
|
||
data = r.get_json() or {}
|
||
self.assertIn("План", data.get("reply", ""))
|
||
self.assertEqual(mock_post.call_count, 2)
|
||
trace = data.get("tool_trace") or []
|
||
self.assertTrue(trace)
|
||
self.assertEqual(trace[0].get("tool"), "slim_sql")
|
||
self.assertTrue(trace[0].get("ok"))
|
||
|
||
@patch("app.services.local_llm_client.requests.post")
|
||
def test_llm_chat_sql_execute_blocked_when_sql_mode_preview(self, mock_post) -> None:
|
||
def _post(url: str, **kwargs: object) -> MagicMock:
|
||
body = kwargs.get("json") or {}
|
||
msgs = body.get("messages") or []
|
||
has_tool = any(
|
||
isinstance(m, dict) and m.get("role") == "tool" for m in msgs
|
||
)
|
||
m = MagicMock()
|
||
m.status_code = 200
|
||
if not has_tool:
|
||
m.json.return_value = {
|
||
"choices": [
|
||
{
|
||
"message": {
|
||
"role": "assistant",
|
||
"tool_calls": [
|
||
{
|
||
"id": "c1",
|
||
"type": "function",
|
||
"function": {
|
||
"name": "slim_sql",
|
||
"arguments": json.dumps(
|
||
{
|
||
"bind": "recipes",
|
||
"sql": "SELECT 1 AS x",
|
||
"mode": "execute",
|
||
}
|
||
),
|
||
},
|
||
}
|
||
],
|
||
}
|
||
}
|
||
]
|
||
}
|
||
else:
|
||
m.json.return_value = {
|
||
"choices": [{"message": {"role": "assistant", "content": "Ок."}}]
|
||
}
|
||
return m
|
||
|
||
mock_post.side_effect = _post
|
||
self._login()
|
||
r = self.client.post(
|
||
"/api/admin/llm/chat",
|
||
data=json.dumps(
|
||
{
|
||
"messages": [{"role": "user", "content": "Данные"}],
|
||
"sql_mode": "preview",
|
||
"tools_enabled": True,
|
||
}
|
||
),
|
||
content_type="application/json",
|
||
)
|
||
self.assertEqual(r.status_code, 200)
|
||
trace = (r.get_json() or {}).get("tool_trace") or []
|
||
self.assertTrue(trace)
|
||
self.assertFalse(trace[0].get("ok"))
|
||
|
||
|
||
class AdminLLMRateLimitTests(unittest.TestCase):
|
||
"""Второй POST llm/summary подряд должен получить 429 при ненулевом интервале."""
|
||
|
||
def setUp(self) -> None:
|
||
reset_llm_rate_limit_for_tests()
|
||
self._prev_security_settings_path = os.environ.get("WESP_SECURITY_SETTINGS_PATH")
|
||
self._prev_network_diagnostics_path = os.environ.get("WESP_NETWORK_DIAGNOSTICS_PATH")
|
||
_sec_dir = tempfile.mkdtemp(prefix="wesp-llm-rl-")
|
||
os.environ["WESP_SECURITY_SETTINGS_PATH"] = os.path.join(_sec_dir, "wesp_security_settings.json")
|
||
_diag_dir = tempfile.mkdtemp(prefix="wesp-llm-rl-d-")
|
||
os.environ["WESP_NETWORK_DIAGNOSTICS_PATH"] = os.path.join(
|
||
_diag_dir, "wesp_network_diagnostics.json"
|
||
)
|
||
|
||
class R(AdminPanelTestConfig):
|
||
WESP_ADMIN_LLM_ENABLED = True
|
||
WESP_ADMIN_LLM_RATE_LIMIT_SEC = 90.0
|
||
|
||
self.app = create_app(R)
|
||
self.client = self.app.test_client()
|
||
self.ctx = self.app.app_context()
|
||
self.ctx.push()
|
||
os.makedirs(R.DATA_DIR, exist_ok=True)
|
||
db.create_all()
|
||
db.session.add(
|
||
WebUser(
|
||
login="admin",
|
||
password_hash=generate_password_hash("admin-secret"),
|
||
is_superuser=True,
|
||
)
|
||
)
|
||
db.session.commit()
|
||
|
||
def tearDown(self) -> None:
|
||
db.session.remove()
|
||
db.drop_all()
|
||
self.ctx.pop()
|
||
reset_llm_rate_limit_for_tests()
|
||
if self._prev_security_settings_path is None:
|
||
os.environ.pop("WESP_SECURITY_SETTINGS_PATH", None)
|
||
else:
|
||
os.environ["WESP_SECURITY_SETTINGS_PATH"] = self._prev_security_settings_path
|
||
if self._prev_network_diagnostics_path is None:
|
||
os.environ.pop("WESP_NETWORK_DIAGNOSTICS_PATH", None)
|
||
else:
|
||
os.environ["WESP_NETWORK_DIAGNOSTICS_PATH"] = self._prev_network_diagnostics_path
|
||
|
||
@patch("app.services.local_llm_client.requests.post")
|
||
def test_second_summary_returns_429(self, mock_post) -> None:
|
||
mock_post.return_value.status_code = 200
|
||
mock_post.return_value.json.return_value = _openai_chat_response("x")
|
||
self.client.post("/api/auth/login", json={"login": "admin", "password": "admin-secret"})
|
||
body = json.dumps({"date_from": "2026-01-01", "date_to": "2026-01-07", "focus": "both"})
|
||
r1 = self.client.post("/api/admin/llm/summary", data=body, content_type="application/json")
|
||
self.assertEqual(r1.status_code, 200)
|
||
r2 = self.client.post("/api/admin/llm/summary", data=body, content_type="application/json")
|
||
self.assertEqual(r2.status_code, 429)
|
||
self.assertIn("Подождите", r2.get_json().get("message", ""))
|
||
self.assertEqual(mock_post.call_count, 1)
|
||
|
||
|
||
class AdminLLMActivityLogTests(unittest.TestCase):
|
||
def test_append_read_and_format(self) -> None:
|
||
tmp = tempfile.mkdtemp(prefix="wesp-llm-log-")
|
||
log_path = os.path.join(tmp, "wesp-llm.jsonl")
|
||
|
||
class _App:
|
||
config = {
|
||
"TESTING": False,
|
||
"WESP_ADMIN_LLM_LOG_PATH": log_path,
|
||
"BASE_DIR": tmp,
|
||
"WESP_ADMIN_LLM_LOG_RETENTION_DAYS": 7,
|
||
}
|
||
|
||
app = _App()
|
||
append_admin_llm_activity(
|
||
app,
|
||
event="chat",
|
||
level="ok",
|
||
user_login="admin",
|
||
detail="ctx=light tools=on",
|
||
duration_ms=120,
|
||
meta={"tool_steps": 2},
|
||
)
|
||
rows, path = read_admin_llm_activity(app, max_lines=10)
|
||
self.assertTrue(os.path.isfile(path))
|
||
self.assertEqual(len(rows), 1)
|
||
self.assertEqual(rows[0].get("event"), "chat")
|
||
self.assertEqual(rows[0].get("user"), "admin")
|
||
text = format_llm_log_line_text(rows[0])
|
||
self.assertIn("llm/chat", text)
|
||
self.assertIn("120 ms", text)
|