Click/load_test.py

"""
Нагрузочный тест: N визитов через BrowserPool.

Использование:
    python load_test.py --n 5
    python load_test.py --n 10 --url https://example.com --time 15 --workers 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
from services.browser_pool import BrowserPool

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


async def submit_visit(
    index: int,
    pool: BrowserPool,
    url: str,
    visit_time: int,
    results: List[BrowserResult],
) -> None:
    t0 = time.perf_counter()
    try:
        result = await pool.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,
        ))
        status = "OK" if result.success else f"FAIL: {result.error}"
        print(f"  [#{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"  [#{index:>2}] EXCEPTION: {e} — {duration:.1f}s")


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

    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 < workers:
        print(f"  ВНИМАНИЕ: прокси {proxy_count} < воркеров {workers}.")
        print(f"  Часть воркеров может ждать свободный прокси.\n")

    browser_service = BrowserService(proxy_manager)
    pool = BrowserPool(size=workers, browser_service=browser_service)
    await pool.start()

    results: List[BrowserResult] = []
    ram_before = _mem_mb()
    print(f"  RAM до старта: {ram_before:.0f} MB")
    print(f"\n  Отправляем {n} визитов в пул ({workers} воркеров)...\n")

    wall_start = time.perf_counter()

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

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

    await pool.stop()
    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) / workers if workers else 0:.0f} MB")

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

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


def main():
    parser = argparse.ArgumentParser(description="Load test через BrowserPool")
    parser.add_argument("--n", type=int, default=5, help="Число визитов (по умолчанию: 5)")
    parser.add_argument("--workers", 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()

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


if __name__ == "__main__":
    main()