84 lines
5.3 KiB
Python
84 lines
5.3 KiB
Python
"""AI-помощник: настройки (ключ не утекает), фильтр тематики, инструменты с правами пользователя, лимиты."""
|
|
import os
|
|
import tempfile
|
|
|
|
os.environ.setdefault("DATABASE_URL", os.environ.get("TEST_DATABASE_URL") or "sqlite:///" + os.path.join(tempfile.mkdtemp(), "as.db").replace("\\", "/"))
|
|
os.environ.setdefault("SCHEDULER_INTERVAL", "3600")
|
|
os.environ.setdefault("SECRET_KEY", "test-secret-key-for-pytest-0123456789abcdef")
|
|
os.environ.setdefault("DATA_DIR", tempfile.mkdtemp())
|
|
|
|
from fastapi.testclient import TestClient # noqa: E402
|
|
|
|
from app import assistant as A # noqa: E402
|
|
from app.main import app # noqa: E402
|
|
|
|
V1 = "/api/v1"
|
|
|
|
|
|
def test_assistant(monkeypatch):
|
|
with TestClient(app) as c:
|
|
adm = {"Authorization": "Bearer " + c.post(f"{V1}/auth/login", json={"username": "admin", "password": "admin"}).json()["access_token"]}
|
|
c.post(f"{V1}/users", headers=adm, json={"username": "as_tech", "password": "Pass1234567", "full_name": "Помощников Тест"})
|
|
tech = {"Authorization": "Bearer " + c.post(f"{V1}/auth/login", json={"username": "as_tech", "password": "Pass1234567"}).json()["access_token"]}
|
|
ask = lambda q, h=tech: c.post(f"{V1}/assistant/chat", headers=h, json={"messages": [{"role": "user", "content": q}]})
|
|
|
|
# выключен по умолчанию
|
|
assert c.get(f"{V1}/assistant/status", headers=tech).json() == {"enabled": False}
|
|
assert ask("где задания?").status_code == 503
|
|
# настройки — только админ; без ключа включить нельзя
|
|
assert c.put(f"{V1}/assistant/config", headers=tech, json={"enabled": True}).status_code == 403
|
|
assert c.put(f"{V1}/assistant/config", headers=adm, json={"enabled": True, "model": "m"}).status_code == 422
|
|
assert c.put(f"{V1}/assistant/config", headers=adm, json={"base_url": "ftp://x"}).status_code == 422
|
|
r = c.put(f"{V1}/assistant/config", headers=adm, json={"enabled": True, "provider": "openai", "model": "test-model",
|
|
"api_key": "sk-secret-value-1234", "base_url": "https://llm.example/v1"})
|
|
assert r.status_code == 200 and r.json()["api_key_set"] and "sk-secret" not in r.text and r.json()["api_key_hint"] == "…1234"
|
|
assert "sk-secret" not in c.get(f"{V1}/settings", headers=adm).text
|
|
assert c.get(f"{V1}/assistant/status", headers=tech).json() == {"enabled": True}
|
|
|
|
calls = []
|
|
|
|
async def fake(cfg, system, messages):
|
|
calls.append((system, messages))
|
|
q = messages[-1]["content"]
|
|
if "стих" in q:
|
|
return "OFFTOPIC"
|
|
if "задани" in q and not q.startswith("РЕЗУЛЬТАТ"):
|
|
return 'CALL my_tasks {}'
|
|
if "склад" in q and not q.startswith("РЕЗУЛЬТАТ"):
|
|
return 'CALL warehouses {}'
|
|
if q.startswith("РЕЗУЛЬТАТ"):
|
|
return "Ответ по данным: " + q[:200]
|
|
return "Откройте [Суточные задания](#/tasks)."
|
|
monkeypatch.setattr(A, "complete", fake)
|
|
|
|
# системный промт знает пользователя и его роль; справка в нём
|
|
r = ask("Как поставить задание?")
|
|
assert r.status_code == 200, r.text
|
|
sys = calls[-1][0]
|
|
assert "Помощников Тест" in sys and "техник" in sys and "Задачи и поручения" in sys and "OFFTOPIC" in sys
|
|
# не по теме — вежливый отказ, а не ответ модели
|
|
r = ask("Напиши стих про весну").json()
|
|
assert r["offtopic"] and r["reply"] == A.REFUSAL
|
|
# инструмент вызывается от имени пользователя: его задания
|
|
c.post(f"{V1}/tasks", headers=tech, json={"title": "Проверить насос помощником"})
|
|
r = ask("Какие у меня задания?").json()
|
|
assert r["tools"] == ["my_tasks"] and "Проверить насос помощником" in r["reply"]
|
|
# нет прав — помощник получает отказ, а не данные
|
|
out = __import__("asyncio").run(A.run_tool(app, tech["Authorization"][7:], "equipment_units", {"id": "abc"}))
|
|
assert "числовой" in out
|
|
# длина вопроса и частота
|
|
assert ask("x" * 1500).status_code == 422
|
|
A._hits.clear()
|
|
monkeypatch.setattr(A, "RATE", (2, 3600))
|
|
assert ask("где наряды?").status_code == 200 and ask("где наряды?").status_code == 200
|
|
assert ask("где наряды?").status_code == 429
|
|
A._hits.clear()
|
|
# ключ не попадает в резервную копию
|
|
b = c.post(f"{V1}/backups", headers=adm, json={}).json()
|
|
import io, json, zipfile
|
|
z = zipfile.ZipFile(io.BytesIO(c.get(f"{V1}/backups/{b['name']}", headers=adm).content))
|
|
assert "sk-secret" not in z.read("db/settings.json").decode()
|
|
# выключение
|
|
c.put(f"{V1}/assistant/config", headers=adm, json={"enabled": False})
|
|
assert c.get(f"{V1}/assistant/status", headers=tech).json() == {"enabled": False}
|