sber_med/test_fpg.py

"""
Тестовый скрипт для SDK ФПГ (GigaDoc).
Покрывает: авторизацию → создание сессии → WebSocket (изображения) → удаление сессии.

Использование:
    python3 test_fpg.py [--ws-host HOST] [--ws-port PORT] [--image-dir DIR]
"""

import argparse
import asyncio
import glob
import json
import os
import sys
import time

import requests

# ─── Конфигурация ─────────────────────────────────────────────
OAUTH_URL = "https://oauth.sbermed.ai/token"
CLIENT_ID = "client_fpg_front"
CLIENT_SECRET = os.environ.get("FPG_CLIENT_SECRET", "")

REST_HOST = "gigadoc-sdk.sbermed.ai"
REST_BASE = f"https://{REST_HOST}"

DEVICE_ID = "web_prod"
FORM_FACTOR = "web_prod"


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

    uri = f"wss://{ws_host}:{ws_port}/api/v2/ws/predict/{uid}/image"
    print(f"[WS] Подключение: {uri}")

    async with websockets.connect(
        uri,
        open_timeout=10,
        additional_headers={"Authorization": f"Bearer {token}"},
    ) 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 ФПГ")
    parser.add_argument("--ws-host", default="sdk.sbermed.ai", help="WebSocket host")
    parser.add_argument("--ws-port", type=int, default=8443, 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)")
    print("=" * 60)

    # 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()