248 lines
9.0 KiB
Python
248 lines
9.0 KiB
Python
#!/usr/bin/env python3
|
||
"""SMC model autotest — проверяет все загруженные модели классификации после обновления.
|
||
|
||
Использование:
|
||
python3 test_smc_models.py
|
||
python3 test_smc_models.py --host 10.101.60.27
|
||
python3 test_smc_models.py --host 10.101.60.27 --port 6181
|
||
python3 test_smc_models.py --asserts smc_asserts.json
|
||
|
||
Файл ассертов (JSON):
|
||
{
|
||
"asserts": [
|
||
{
|
||
"model": "МФЦ НСК",
|
||
"text": "записаться на приём", // опционально, свой текст
|
||
"checks": {
|
||
"min_found": 1, // мин. количество классов
|
||
"text_contains": "Запись", // подстрока в top классе
|
||
"confidence_min": 0.5 // мин. confidence топ-класса
|
||
}
|
||
}
|
||
]
|
||
}
|
||
|
||
Возвращает exit code 0 если все модели прошли, иначе 1.
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
import os
|
||
import unittest
|
||
from urllib.parse import quote
|
||
|
||
try:
|
||
import urllib.request as urlreq
|
||
import urllib.error as urlerr
|
||
except ImportError:
|
||
urlreq = None
|
||
|
||
GREEN = '\033[92m'
|
||
RED = '\033[91m'
|
||
YELLOW = '\033[93m'
|
||
BOLD = '\033[1m'
|
||
RESET = '\033[0m'
|
||
|
||
def _disable_colors():
|
||
global GREEN, RED, YELLOW, BOLD, RESET
|
||
GREEN = RED = YELLOW = BOLD = RESET = ''
|
||
|
||
if not os.isatty(sys.stdout.fileno()) or os.environ.get('NO_COLOR'):
|
||
_disable_colors()
|
||
|
||
DEFAULT_HOST = "10.101.60.27"
|
||
DEFAULT_PORT = 6181
|
||
BASE_PATH = "/smc"
|
||
|
||
TEXT_RULES = [
|
||
("филиал", "родники"),
|
||
("offices", "родники"),
|
||
("мфц", "записаться в мфц"),
|
||
("mfc", "записаться в мфц"),
|
||
("ekc", "записаться в екц"),
|
||
("екц", "записаться в екц"),
|
||
("minsoc", "какие выплаты"),
|
||
("минсоц", "какие выплаты"),
|
||
("техническая", "подключите оператора"),
|
||
("поддержка", "подключите оператора"),
|
||
("тест", "тестовый запрос"),
|
||
("test", "test query"),
|
||
]
|
||
|
||
|
||
def pick_text(model_id):
|
||
name_lower = model_id.lower()
|
||
for keyword, text in TEXT_RULES:
|
||
if keyword in name_lower:
|
||
return text
|
||
return "Здравствуйте, хочу записаться на приём"
|
||
|
||
|
||
def load_asserts(path):
|
||
if not path:
|
||
return []
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
return data.get("asserts", [])
|
||
|
||
|
||
class SmcAPI:
|
||
def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT):
|
||
self.base = f"http://{host}:{port}{BASE_PATH}"
|
||
|
||
def get_models(self):
|
||
return json.loads(self._request(f"{self.base}/models"))
|
||
|
||
def classify(self, model_id, text):
|
||
encoded = quote(model_id, safe="")
|
||
url = f"{self.base}/classify/{encoded}?text={quote(text)}"
|
||
return json.loads(self._request(url))
|
||
|
||
def _request(self, url, timeout=30):
|
||
try:
|
||
with urlreq.urlopen(url, timeout=timeout) as r:
|
||
return r.read().decode("utf-8")
|
||
except urlerr.HTTPError as e:
|
||
raise RuntimeError(f"HTTP {e.code}: {e.read().decode('utf-8')[:200]}")
|
||
except urlerr.URLError as e:
|
||
raise RuntimeError(f"Connection failed: {e.reason}")
|
||
|
||
|
||
class TestSmcModels(unittest.TestCase):
|
||
@classmethod
|
||
def setUpClass(cls):
|
||
cls.api = SmcAPI(cls.args.host, cls.args.port)
|
||
cls.asserts = load_asserts(cls.args.asserts)
|
||
data = cls.api.get_models()
|
||
TestSmcModels.models = data
|
||
models = data.get("models", [])
|
||
handlers = data.get("handlers", [])
|
||
print(f"\n models: {len(models)}, handlers: {len(handlers)}")
|
||
if cls.asserts:
|
||
print(f" asserts loaded: {len(cls.asserts)}")
|
||
if len(models) + len(handlers) == 0:
|
||
raise unittest.SkipTest("Нет ни одной модели")
|
||
|
||
def _find_assert(self, model_id):
|
||
for a in self.asserts:
|
||
if a["model"] == model_id:
|
||
return a
|
||
for a in self.asserts:
|
||
if a["model"] in model_id:
|
||
return a
|
||
return None
|
||
|
||
def _check_assert(self, model_id, result):
|
||
a = self._find_assert(model_id)
|
||
if not a:
|
||
return
|
||
checks = a.get("checks", {})
|
||
classes = result.get("classes", [])
|
||
if "min_found" in checks:
|
||
self.assertGreaterEqual(len(classes), checks["min_found"],
|
||
f"{model_id}: expected >= {checks['min_found']} classes, got {len(classes)}")
|
||
if classes and "confidence_min" in checks:
|
||
self.assertGreaterEqual(classes[0].get("confidence", 0), checks["confidence_min"],
|
||
f"{model_id}: top confidence < {checks['confidence_min']}")
|
||
|
||
def test_01_list_models(self):
|
||
"""Получение списка моделей"""
|
||
data = self.models
|
||
self.assertIn("models", data)
|
||
self.assertIn("handlers", data)
|
||
print(f"\n models: {len(data['models'])}")
|
||
print(f" handlers: {len(data['handlers'])}")
|
||
|
||
def _fmt_result(self, result):
|
||
classes = result.get("classes", [])
|
||
if classes:
|
||
top = classes[0]
|
||
return f"{top.get('class','?')} ({top.get('confidence',0):.3f})"
|
||
return "∅"
|
||
|
||
def _run_section(self, section_name, model_ids):
|
||
results = []
|
||
for mid in model_ids:
|
||
if not mid:
|
||
continue
|
||
text = pick_text(mid)
|
||
a = self._find_assert(mid)
|
||
if a and "text" in a:
|
||
text = a["text"]
|
||
try:
|
||
result = self.api.classify(mid, text)
|
||
classes = result.get("classes", [])
|
||
if not classes:
|
||
results.append(("empty", mid, None, None))
|
||
else:
|
||
self._check_assert(mid, result)
|
||
results.append(("ok", mid, result, None))
|
||
except AssertionError:
|
||
results.append(("assert_fail", mid, None, None))
|
||
raise
|
||
except Exception as e:
|
||
results.append(("error", mid, None, str(e)))
|
||
|
||
passed = sum(1 for r in results if r[0] == "ok")
|
||
empty = sum(1 for r in results if r[0] == "empty")
|
||
errors = [r for r in results if r[0] == "error"]
|
||
|
||
print(f"\n{BOLD}── {section_name} ──────────────────────────────────────────────{RESET}")
|
||
for status, mid, result, extra in results:
|
||
if status == "ok" and self.args.verbose:
|
||
print(f" {GREEN}✓ {mid}{RESET} → {self._fmt_result(result)}")
|
||
elif status == "ok":
|
||
print(f" {GREEN}✓ {mid}{RESET}")
|
||
elif status == "empty":
|
||
print(f" {RED}✗ {mid}{RESET} → empty")
|
||
elif status == "assert_fail":
|
||
print(f" {RED}✗ {mid}{RESET} → assert fail")
|
||
elif status == "error":
|
||
print(f" {YELLOW}⚠ {mid}{RESET} → {extra}")
|
||
color = GREEN if empty == 0 and not errors else RED
|
||
print(f" {BOLD}{color}─── {passed} passed, {empty} empty{RESET}", end="")
|
||
if errors:
|
||
print(f", {len(errors)} errors", end="")
|
||
print()
|
||
|
||
error_msgs = [f" {r[1]}: {r[3]}" for r in errors]
|
||
if error_msgs:
|
||
self.fail(f"Ошибки {section_name}:\n" + "\n".join(error_msgs))
|
||
|
||
def test_02_models(self):
|
||
"""Классификация через модели"""
|
||
models = self.models.get("models", [])
|
||
if not models:
|
||
self.skipTest("Нет моделей")
|
||
self._run_section("MODELS", models)
|
||
|
||
def test_03_handlers(self):
|
||
"""Классификация через handler модели"""
|
||
handlers = self.models.get("handlers", [])
|
||
if not handlers:
|
||
self.skipTest("Нет handler моделей")
|
||
self._run_section("HANDLERS", handlers)
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="SMC model autotest")
|
||
parser.add_argument("--host", default=DEFAULT_HOST, help=f"Хост (default: {DEFAULT_HOST})")
|
||
parser.add_argument("--port", type=int, default=DEFAULT_PORT, help=f"Порт (default: {DEFAULT_PORT})")
|
||
parser.add_argument("--asserts", default=None, help="JSON файл с ассертами")
|
||
parser.add_argument("--verbose", "-v", action="store_true", help="Показать результат каждой модели")
|
||
parser.add_argument("--no-color", action="store_true", help="Отключить цвета в выводе")
|
||
args = parser.parse_args()
|
||
if args.no_color:
|
||
_disable_colors()
|
||
TestSmcModels.args = args
|
||
|
||
suite = unittest.TestLoader().loadTestsFromTestCase(TestSmcModels)
|
||
runner = unittest.TextTestRunner(verbosity=2)
|
||
result = runner.run(suite)
|
||
sys.exit(0 if result.wasSuccessful() else 1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|