/pentestgpt_legacy/test_connection.py

"""Per-provider connectivity check.

A lightweight wrapper over the smoke test that probes the first current-tier
model of each configured provider — handy for quickly confirming your keys work
before starting a session. For full coverage use ``pentestgpt-legacy --smoke-test``.
"""

from __future__ import annotations

import asyncio

from pentestgpt_legacy.config import configured_providers
from pentestgpt_legacy.llm.registry import ALL_SPECS, PROVIDERS
from pentestgpt_legacy.smoke_test import PASS, _probe


def _flagship_per_provider() -> list:
    ready = set(configured_providers())
    chosen = []
    seen: set[str] = set()
    for spec in ALL_SPECS:
        if spec.provider in ready and spec.provider not in seen and not spec.legacy:
            chosen.append(spec)
            seen.add(spec.provider)
    return chosen


async def _run() -> bool:
    specs = _flagship_per_provider()
    if not specs:
        print("No providers configured. Set an API key in your environment or .env file.")
        return False

    print("\n=== PentestGPT connectivity test ===\n")
    all_ok = True
    for spec in specs:
        status, detail = await _probe(spec)
        mark = "✓" if status == PASS else "✗"
        print(f"  {mark} {PROVIDERS[spec.provider].label:<16} {spec.id:<28} {detail}")
        all_ok = all_ok and status == PASS
    print()
    return all_ok


def test_connection() -> bool:
    """Synchronous entry point; returns True if all configured providers connect."""
    return asyncio.run(_run())


if __name__ == "__main__":
    raise SystemExit(0 if test_connection() else 1)