227 lines
8.3 KiB
Python
227 lines
8.3 KiB
Python
"""
|
||
Тестовый скрипт для SDK ФПГ (GigaDoc, персональный стенд DNR).
|
||
Покрывает: авторизацию → создание сессии → WebSocket (изображения) → удаление сессии.
|
||
|
||
Использование:
|
||
python3 test_fpg.py [--ws-host HOST] [--ws-port PORT] [--image-dir DIR]
|
||
|
||
Хосты можно переопределить через env: FPG_REST_HOST, FPG_WS_HOST, FPG_WS_PORT.
|
||
Хост может содержать путь (domain.com/v1/giga) — он учитывается при сборке URL.
|
||
"""
|
||
|
||
import argparse
|
||
import asyncio
|
||
import glob
|
||
import json
|
||
import os
|
||
import sys
|
||
import time
|
||
|
||
import requests
|
||
|
||
# ─── Конфигурация ─────────────────────────────────────────────
|
||
OAUTH_URL = os.environ.get("FPG_OAUTH_URL", "https://oauth.sbermed.ai/token")
|
||
CLIENT_ID = os.environ.get("FPG_CLIENT_ID", "client_fpg_front_dnr")
|
||
CLIENT_SECRET = os.environ.get("FPG_CLIENT_SECRET", "")
|
||
|
||
REST_HOST = os.environ.get("FPG_REST_HOST", "gigadoc-dnr-sdk.sbermed.ai")
|
||
WS_HOST = os.environ.get("FPG_WS_HOST", "gigadoc-dnr-sdk.sbermed.ai")
|
||
WS_PORT = int(os.environ.get("FPG_WS_PORT", "8443"))
|
||
WS_ORIGIN = os.environ.get("FPG_WS_ORIGIN", "https://gigadoc-dnr-sdk.sbermed.ai")
|
||
|
||
DEVICE_ID = "web_prod"
|
||
FORM_FACTOR = "web_prod"
|
||
|
||
|
||
def split_host_path(host: str) -> tuple:
|
||
"""Разделить 'host/path' на (scheme, hostname, base_path)."""
|
||
host = host.strip().rstrip("/")
|
||
scheme = ""
|
||
if "://" in host:
|
||
scheme, host = host.split("://", 1)
|
||
parts = host.split("/", 1)
|
||
hostname = parts[0]
|
||
base_path = "/" + parts[1] if len(parts) > 1 else ""
|
||
return scheme, hostname, base_path
|
||
|
||
|
||
_rest_scheme, _rest_hostname, _rest_base_path = split_host_path(REST_HOST)
|
||
REST_BASE = f"{_rest_scheme or 'https'}://{_rest_hostname}{_rest_base_path}"
|
||
|
||
_ws_scheme, WS_HOSTNAME, WS_BASE_PATH = split_host_path(WS_HOST)
|
||
|
||
|
||
def get_token() -> str:
|
||
"""Получить access token через client_credentials."""
|
||
resp = requests.post(
|
||
OAUTH_URL,
|
||
data={
|
||
"grant_type": "client_credentials",
|
||
"client_id": CLIENT_ID,
|
||
"client_secret": CLIENT_SECRET,
|
||
},
|
||
timeout=10,
|
||
)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
print(f"[AUTH] Токен получен (действует {data['expires_in']} сек)")
|
||
return data["access_token"]
|
||
|
||
|
||
def create_session(token: str) -> str:
|
||
"""Создать сессию и вернуть uid."""
|
||
resp = requests.post(
|
||
f"{REST_BASE}/api/v1/predict/new",
|
||
params={"device_id": DEVICE_ID, "form_factor": FORM_FACTOR},
|
||
headers={
|
||
"Authorization": f"Bearer {token}",
|
||
"Accept": "application/json",
|
||
},
|
||
timeout=10,
|
||
)
|
||
resp.raise_for_status()
|
||
uid = resp.json()["uid"]
|
||
print(f"[SESSION] Сессия создана: {uid}")
|
||
return uid
|
||
|
||
|
||
def delete_session(token: str, uid: str):
|
||
"""Удалить сессию."""
|
||
resp = requests.delete(
|
||
f"{REST_BASE}/api/v1/predict/{uid}",
|
||
headers={"Authorization": f"Bearer {token}"},
|
||
timeout=10,
|
||
)
|
||
print(f"[SESSION] Сессия {uid} удалена ({resp.status_code})")
|
||
|
||
|
||
async def ws_scan(token: str, uid: str, ws_host: str, ws_port: int, image_dir: str):
|
||
"""Подключиться по WebSocket, отправить изображения, получить результат."""
|
||
import websockets
|
||
|
||
# ws_host может быть с путём (domain.com/v1/giga) — разделяем.
|
||
_scheme, _hostname, _base_path = split_host_path(ws_host)
|
||
_base_path = _base_path or WS_BASE_PATH
|
||
|
||
uri = f"wss://{_hostname}:{ws_port}{_base_path}/api/v2/ws/predict/{uid}/image"
|
||
print(f"[WS] Подключение: {uri}")
|
||
|
||
# Авторизация через subprotocol bearer.<token> (как фронт),
|
||
# Origin — обязательный заголовок.
|
||
async with websockets.connect(
|
||
uri,
|
||
open_timeout=10,
|
||
subprotocols=[f"bearer.{token}"],
|
||
extra_headers={"Origin": WS_ORIGIN},
|
||
) as ws:
|
||
print("[WS] Соединение установлено")
|
||
|
||
# Собираем изображения
|
||
patterns = ["*.jpg", "*.jpeg", "*.png"]
|
||
images = []
|
||
for p in patterns:
|
||
images.extend(glob.glob(os.path.join(image_dir, p)))
|
||
images.sort()
|
||
|
||
if not images:
|
||
print(f"[WS] Нет изображений в {image_dir}, отправляю тестовое изображение")
|
||
images = ["/tmp/test_face.jpg"]
|
||
|
||
print(f"[WS] Отправляю {len(images)} изображений...")
|
||
|
||
for i, img_path in enumerate(images):
|
||
with open(img_path, "rb") as f:
|
||
data = f.read()
|
||
await ws.send(data)
|
||
|
||
try:
|
||
resp = await asyncio.wait_for(ws.recv(), timeout=5)
|
||
result = json.loads(resp)
|
||
progress = result.get("progress", "?")
|
||
|
||
# Краткий вывод прогресса
|
||
if i % 30 == 0 or progress == 1.0:
|
||
print(
|
||
f" [{i+1}/{len(images)}] progress={progress} "
|
||
f"face_ok={result.get('face_ok')}"
|
||
)
|
||
|
||
# Финальный результат
|
||
if progress == 1.0 or (isinstance(progress, (int, float)) and progress >= 1.0):
|
||
print("\n" + "=" * 60)
|
||
print("📊 ФИНАЛЬНЫЙ РЕЗУЛЬТАТ ДИАГНОСТИКИ:")
|
||
print("=" * 60)
|
||
print(json.dumps(result, indent=2, ensure_ascii=False))
|
||
return result
|
||
|
||
except asyncio.TimeoutError:
|
||
if i % 50 == 0:
|
||
print(f" [{i+1}/{len(images)}] ожидание ответа...")
|
||
|
||
print("[WS] Соединение закрыто без финального результата")
|
||
return None
|
||
|
||
|
||
def create_test_image():
|
||
"""Создать тестовое JPEG-изображение если нет своих."""
|
||
try:
|
||
from PIL import Image
|
||
|
||
img = Image.new("RGB", (1280, 720), color=(128, 128, 128))
|
||
img.save("/tmp/test_face.jpg", "JPEG", quality=80)
|
||
print("[INIT] Создано тестовое изображение /tmp/test_face.jpg")
|
||
except ImportError:
|
||
print("[INIT] PIL не установлен, тестовое изображение не создано")
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Тест SDK ФПГ (DNR)")
|
||
parser.add_argument("--ws-host", default=WS_HOSTNAME, help="WebSocket host (может быть с путём)")
|
||
parser.add_argument("--ws-port", type=int, default=WS_PORT, help="WebSocket port")
|
||
parser.add_argument("--image-dir", default=".", help="Директория с JPEG-файлами")
|
||
parser.add_argument("--rest-only", action="store_true", help="Только REST API (без WS)")
|
||
args = parser.parse_args()
|
||
|
||
print("=" * 60)
|
||
print("🏥 Тестирование SDK ФПГ (GigaDoc DNR)")
|
||
print("=" * 60)
|
||
print(f"[CONFIG] REST: {REST_BASE}")
|
||
print(f"[CONFIG] WS: wss://{args.ws_host}:{args.ws_port}")
|
||
|
||
# 0. Тестовое изображение
|
||
create_test_image()
|
||
|
||
# 1. Авторизация
|
||
token = get_token()
|
||
|
||
# 2. Создание сессии
|
||
uid = create_session(token)
|
||
|
||
# 3. WebSocket
|
||
if not args.rest_only:
|
||
try:
|
||
result = asyncio.run(
|
||
ws_scan(token, uid, args.ws_host, args.ws_port, args.image_dir)
|
||
)
|
||
except Exception as e:
|
||
print(f"\n[WS] Ошибка: {type(e).__name__}: {e}")
|
||
print("[WS] Проверьте доступность WebSocket хоста")
|
||
result = None
|
||
else:
|
||
result = None
|
||
print("[SKIP] WebSocket тест пропущен (--rest-only)")
|
||
|
||
# 4. Удаление сессии
|
||
delete_session(token, uid)
|
||
|
||
# Итог
|
||
print("\n" + "=" * 60)
|
||
if result:
|
||
print("✅ Тест завершён успешно — результат получен")
|
||
else:
|
||
print("⚠️ REST API работает, WebSocket недоступен")
|
||
print("=" * 60)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |