Click/load_test.py

"""
Нагрузочный тест: N параллельных браузеров на X секунд.

Использование:
    python load_test.py --n 5
    python load_test.py --n 10 --url https://example.com --time 15
    python load_test.py --n 10 --parallel 3   # не более 3 одновременных запусков
"""

import asyncio
import sys
import time
import argparse
from pathlib import Path
from dataclasses import dataclass
from typing import List, Optional

sys.path.insert(0, str(Path(__file__).parent))

import psutil

from core.logger import setup_logger
from managers.proxy_manager import ProxyManager
from services.browser_service import BrowserService

logger = setup_logger(name="load_test", log_level="INFO", log_file="logs/load_test.log")

PROCESS = psutil.Process()


def _mem_mb() -> float:
    return PROCESS.memory_info().rss / (1024 * 1024)


@dataclass
class BrowserResult:
    index: int
    success: bool
    duration: float
    error: Optional[str] = None
    final_url: Optional[str] = None


async def run_one_browser(
    index: int,
    browser_service: BrowserService,
    url: str,
    visit_time: int,
    results: List[BrowserResult],
    launch_sem: asyncio.Semaphore,
) -> None:
    # Семафор ограничивает сколько браузеров стартует одновременно.
    # Когда браузер поднялся (AsyncCamoufox.__aenter__ завершён), семафор освобождается
    # и следующий браузер может начать запуск. Так избегаем race на ресурсы ОС.
    async with launch_sem:
        t0 = time.perf_counter()
        try:
            result = await browser_service.visit_page(url, reading_time=visit_time)
            duration = time.perf_counter() - t0
            results.append(BrowserResult(
                index=index,
                success=result.success,
                duration=duration,
                error=result.error,
                final_url=result.final_url,
            ))
            status = "OK" if result.success else f"FAIL: {result.error}"
            print(f"  [Browser {index:>2}] {status} — {duration:.1f}s | RAM: {_mem_mb():.0f} MB")
        except Exception as e:
            duration = time.perf_counter() - t0
            results.append(BrowserResult(index=index, success=False, duration=duration, error=str(e)))
            print(f"  [Browser {index:>2}] EXCEPTION: {e} — {duration:.1f}s")


async def run_load_test(n: int, url: str, visit_time: int, parallel: int) -> None:
    print(f"\n{'='*60}")
    print(f"  Нагрузочный тест")
    print(f"  Браузеров: {n}  |  Одновременный запуск: {parallel}  |  URL: {url}  |  Время: {visit_time}s")

    # Инфо о системе
    mem = psutil.virtual_memory()
    print(f"  RAM всего: {mem.total / (1024**3):.1f} GB  |  Свободно: {mem.available / (1024**3):.1f} GB")
    print(f"  CPU ядер: {psutil.cpu_count()}")
    print(f"{'='*60}\n")

    proxy_manager = ProxyManager()
    await proxy_manager.detect_types()

    proxy_count = proxy_manager.count
    if proxy_count < n:
        print(f"  ВНИМАНИЕ: прокси {proxy_count} < браузеров {n}.")
        print(f"  Некоторые браузеры могут не получить прокси.\n")

    browser_service = BrowserService(proxy_manager)
    results: List[BrowserResult] = []

    ram_before = _mem_mb()
    print(f"  RAM до старта: {ram_before:.0f} MB")
    print(f"\n  Запускаем {n} браузеров (волнами по {parallel})...\n")

    launch_sem = asyncio.Semaphore(parallel)
    wall_start = time.perf_counter()

    tasks = [
        asyncio.create_task(
            run_one_browser(i + 1, browser_service, url, visit_time, results, launch_sem)
        )
        for i in range(n)
    ]
    await asyncio.gather(*tasks, return_exceptions=True)

    wall_elapsed = time.perf_counter() - wall_start
    ram_after = _mem_mb()

    await browser_service.cleanup()

    # --- Итоги ---
    ok = sum(1 for r in results if r.success)
    fail = len(results) - ok
    avg_dur = sum(r.duration for r in results) / len(results) if results else 0
    max_dur = max((r.duration for r in results), default=0)
    min_dur = min((r.duration for r in results), default=0)

    print(f"\n{'='*60}")
    print(f"  ИТОГИ")
    print(f"{'='*60}")
    print(f"  Всего браузеров : {n}")
    print(f"  Успешно         : {ok}")
    print(f"  Ошибки          : {fail}")
    print(f"  Общее время     : {wall_elapsed:.1f}s")
    print(f"  Время посещения : avg={avg_dur:.1f}s  min={min_dur:.1f}s  max={max_dur:.1f}s")
    print(f"  RAM до/после    : {ram_before:.0f} MB → {ram_after:.0f} MB  (Δ {ram_after - ram_before:+.0f} MB)")
    print(f"  RAM на браузер  : ~{(ram_after - ram_before) / n if n else 0:.0f} MB")

    if fail:
        print(f"\n  Ошибки:")
        for r in results:
            if not r.success:
                print(f"    Browser {r.index}: {r.error}")

    print(f"{'='*60}\n")


def main():
    parser = argparse.ArgumentParser(description="Load test: N параллельных браузеров")
    parser.add_argument("--n", type=int, default=3, help="Всего браузеров (по умолчанию: 3)")
    parser.add_argument("--parallel", type=int, default=2, help="Макс. одновременных запусков (по умолчанию: 2)")
    parser.add_argument("--url", type=str, default="https://httpbin.org/headers", help="URL для посещения")
    parser.add_argument("--time", type=int, default=10, dest="visit_time", help="Время посещения в секундах (по умолчанию: 10)")
    args = parser.parse_args()

    parallel = min(args.parallel, args.n)
    try:
        asyncio.run(run_load_test(args.n, args.url, args.visit_time, parallel))
    except KeyboardInterrupt:
        print("\nПрервано.")


if __name__ == "__main__":
    main()