hermes-hub/tests/test_codex_opencode_wizard.py

"""Comprehensive tests for OpenAI Codex OAuth, Multi-Account Isolation, OpenCode Go, and Wizard Clipboard UX.

TEST A: CodexOAuthSession standard device flow (user code -> poll -> token exchange -> save profile auth)
TEST B: CodexOAuthSession manual token / JSON credential fallback
TEST C: Multi-account isolation: multiple Codex profiles saved and loaded independently without overwriting
TEST D: CodexAdapter._resolve_token resolves tokens for individual profiles from ProfileAuthManager
TEST E: CodexAdapter multi-profile switching and failover isolation
TEST F: ProfileAuthManager.get_profile_status works for both Codex OAuth and API Key modes
TEST G: ProfileAuthManager.extract_jwt_identity extracts email from OpenAI/Google JWT payloads
TEST H: HubEntry and enable_clipboard_shortcuts (Ctrl+V, Ctrl+C, Ctrl+X, Ctrl+A, Shift+Insert, Cyrillic shortcuts)
TEST I: Add Account Wizard: OpenCode Go API key paste with whitespace / newline stripping
TEST J: Add Account Wizard: Password masking does not corrupt backing value or clipboard insertion
TEST K: Zero-secret logging across Codex OAuth and API key flows
TEST L: Single completion lock prevents double token exchange / double file writes in CodexOAuthSession
TEST M: Codex OAuth session cleanup on cancellation / modal destroy
"""
from __future__ import annotations

import base64
import json
import os
import sys
import tempfile
import time
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

from antigravity_provider.router.profile_manager import ProfileAuthManager, mask_email
from antigravity_provider.router.router_config import RouterProfileConfig
from antigravity_provider.router.adapters.codex_adapter import CodexAdapter
from antigravity_provider.router.codex_oauth import (
    CodexOAuthSession,
    start_codex_oauth,
    get_codex_oauth_session,
    cancel_codex_oauth_session,
)


@pytest.fixture(autouse=True)
def isolated_hermes_env(tmp_path, monkeypatch):
    """Ensure all tests run in an isolated HERMES_HOME."""
    hermes_dir = tmp_path / "hermes_test"
    hermes_dir.mkdir(parents=True, exist_ok=True)
    monkeypatch.setenv("HERMES_HOME", str(hermes_dir))
    monkeypatch.setenv("LOCALAPPDATA", str(tmp_path / "localappdata"))
    yield hermes_dir


def _make_jwt(payload: dict) -> str:
    """Helper to generate an unsigned test JWT."""
    header_b64 = base64.urlsafe_b64encode(b'{"alg":"none","typ":"JWT"}').decode("utf-8").rstrip("=")
    payload_b64 = base64.urlsafe_b64encode(json.dumps(payload).encode("utf-8")).decode("utf-8").rstrip("=")
    return f"{header_b64}.{payload_b64}.fake_signature"


# ==============================================================================
# TEST A: CodexOAuthSession standard device flow
# ==============================================================================
def test_a_codex_oauth_standard_device_flow(monkeypatch):
    """Test start -> poll -> token exchange -> profile saving."""
    session = CodexOAuthSession("codex-worker-1")

    # Mock usercode endpoint
    monkeypatch.setattr(
        "antigravity_provider.router.codex_oauth._post_json",
        lambda url, payload, **kwargs: {
            "user_code": "TEST-1234",
            "device_auth_id": "dev-auth-xyz",
            "interval": 1,
        } if "usercode" in url else {
            "authorization_code": "auth-code-999",
            "code_verifier": "verifier-111",
        },
    )

    test_id_token = _make_jwt({"email": "dev1@openai.com", "sub": "usr_dev1"})
    monkeypatch.setattr(
        "antigravity_provider.router.codex_oauth._post_form_json",
        lambda url, data, **kwargs: {
            "access_token": "oa-acc-token-123",
            "refresh_token": "oa-ref-token-456",
            "id_token": test_id_token,
        },
    )

    url, code = session.start()
    assert "auth.openai.com/codex/device" in url
    assert code == "TEST-1234"

    # Wait for poll thread to exchange tokens
    deadline = time.time() + 2.5
    while time.time() < deadline and session.status != "completed":
        time.sleep(0.1)

    assert session.status == "completed"
    assert session.completed_profile_info is not None
    assert session.completed_profile_info["email"] == "dev1@openai.com"

    # Verify saved on disk
    auth_data = ProfileAuthManager.load_profile_auth("openai-codex", "codex-worker-1")
    assert auth_data is not None
    assert auth_data["profile_id"] == "codex-worker-1"
    assert auth_data["token"]["access_token"] == "oa-acc-token-123"
    assert auth_data["email"] == "dev1@openai.com"


# ==============================================================================
# TEST B: CodexOAuthSession manual token / JSON credential fallback
# ==============================================================================
def test_b_codex_oauth_manual_token_fallback():
    """Test manual JSON credential insertion into Codex OAuth session."""
    session = CodexOAuthSession("codex-worker-2")
    test_id_token = _make_jwt({"email": "worker2@openai.com", "sub": "usr_w2"})

    raw_json = json.dumps({
        "access_token": "manual-access-token-777",
        "refresh_token": "manual-refresh-token-888",
        "id_token": test_id_token,
    })

    ok, msg = session.handle_manual_input(raw_json)
    assert ok is True
    assert session.status == "completed"

    auth_data = ProfileAuthManager.load_profile_auth("openai-codex", "codex-worker-2")
    assert auth_data is not None
    assert auth_data["profile_id"] == "codex-worker-2"
    assert auth_data["token"]["access_token"] == "manual-access-token-777"
    assert auth_data["email"] == "worker2@openai.com"


# ==============================================================================
# TEST C: Multi-account isolation
# ==============================================================================
def test_c_multi_account_codex_isolation():
    """Verify multiple Codex accounts are saved and loaded independently."""
    # Save account 1 (codex-orch)
    orch_data = {
        "provider": "openai-codex",
        "profile_id": "codex-orch",
        "auth_mode": "oauth",
        "token": {"access_token": "orch-token-1"},
        "email": "orch@company.com",
    }
    ProfileAuthManager.save_profile_auth("openai-codex", "codex-orch", orch_data)

    # Save account 2 (codex-worker-1)
    w1_data = {
        "provider": "openai-codex",
        "profile_id": "codex-worker-1",
        "auth_mode": "oauth",
        "token": {"access_token": "worker1-token-2"},
        "email": "coder@company.com",
    }
    ProfileAuthManager.save_profile_auth("openai-codex", "codex-worker-1", w1_data)

    # Save account 3 (codex-worker-2 as API key)
    w2_data = {
        "provider": "openai-codex",
        "profile_id": "codex-worker-2",
        "auth_mode": "api_key",
        "api_key": "sk-proj-testkey333333333333",
    }
    ProfileAuthManager.save_profile_auth("openai-codex", "codex-worker-2", w2_data)

    # Verify each loads its own independent data
    loaded_orch = ProfileAuthManager.load_profile_auth("openai-codex", "codex-orch")
    loaded_w1 = ProfileAuthManager.load_profile_auth("openai-codex", "codex-worker-1")
    loaded_w2 = ProfileAuthManager.load_profile_auth("openai-codex", "codex-worker-2")

    assert loaded_orch["token"]["access_token"] == "orch-token-1"
    assert loaded_orch["email"] == "orch@company.com"

    assert loaded_w1["token"]["access_token"] == "worker1-token-2"
    assert loaded_w1["email"] == "coder@company.com"

    assert loaded_w2["api_key"] == "sk-proj-testkey333333333333"


# ==============================================================================
# TEST D: CodexAdapter._resolve_token per-profile isolation
# ==============================================================================
def test_d_codex_adapter_resolve_token():
    """Verify CodexAdapter._resolve_token looks up the exact profile credentials from ProfileAuthManager."""
    adapter = CodexAdapter()

    # Create 2 separate profiles
    p1 = RouterProfileConfig(profile_id="codex-orch", provider="openai-codex")
    p2 = RouterProfileConfig(profile_id="codex-worker-1", provider="openai-codex")

    ProfileAuthManager.save_profile_auth("openai-codex", "codex-orch", {
        "provider": "openai-codex", "profile_id": "codex-orch", "token": {"access_token": "token-orch-99"}
    })
    ProfileAuthManager.save_profile_auth("openai-codex", "codex-worker-1", {
        "provider": "openai-codex", "profile_id": "codex-worker-1", "token": {"access_token": "token-worker-11"}
    })

    t1 = adapter._resolve_token(p1)
    t2 = adapter._resolve_token(p2)

    assert t1 == "token-orch-99"
    assert t2 == "token-worker-11"
    assert t1 != t2


# ==============================================================================
# TEST E: CodexAdapter multi-profile switching and failover isolation
# ==============================================================================
def test_e_codex_adapter_switching_and_failover():
    """Test switching profiles in CodexAdapter during execution."""
    adapter = CodexAdapter()

    p_primary = RouterProfileConfig(profile_id="codex-orch", provider="openai-codex")
    p_spare = RouterProfileConfig(profile_id="codex-spare-1", provider="openai-codex")

    ProfileAuthManager.save_profile_auth("openai-codex", "codex-orch", {
        "provider": "openai-codex", "profile_id": "codex-orch", "api_key": "sk-primary-orch-key111"
    })
    ProfileAuthManager.save_profile_auth("openai-codex", "codex-spare-1", {
        "provider": "openai-codex", "profile_id": "codex-spare-1", "api_key": "sk-spare-orch-key222"
    })

    assert adapter._resolve_token(p_primary) == "sk-primary-orch-key111"
    assert adapter._resolve_token(p_spare) == "sk-spare-orch-key222"


# ==============================================================================
# TEST F: ProfileAuthManager.get_profile_status for OAuth and API Key
# ==============================================================================
def test_f_profile_auth_manager_status():
    """Verify get_profile_status properly reports OAuth vs API Key profiles."""
    test_id_token = _make_jwt({"email": "testuser@openai.com", "sub": "usr_1"})
    ProfileAuthManager.save_profile_auth("openai-codex", "codex-orch", {
        "provider": "openai-codex",
        "profile_id": "codex-orch",
        "token": {"access_token": "test-access", "id_token": test_id_token},
    })
    ProfileAuthManager.save_profile_auth("openai-codex", "codex-worker-1", {
        "provider": "openai-codex",
        "profile_id": "codex-worker-1",
        "api_key": "sk-1234567890abcdef1234",
    })

    status_oauth = ProfileAuthManager.get_profile_status("openai-codex", "codex-orch")
    status_key = ProfileAuthManager.get_profile_status("openai-codex", "codex-worker-1")

    assert status_oauth["authenticated"] is True
    assert status_oauth["auth_mode"] == "oauth"
    assert "test***@openai.com" in status_oauth["email_masked"]

    assert status_key["authenticated"] is True
    assert status_key["auth_mode"] == "api_key"
    assert status_key["account_id_masked"] == "sk-...1234"


# ==============================================================================
# TEST G: ProfileAuthManager.extract_jwt_identity
# ==============================================================================
def test_g_extract_jwt_identity():
    """Test extracting email and sub from various standard and custom JWT claims."""
    jwt_standard = _make_jwt({"email": "standard@gmail.com", "sub": "google-sub-123"})
    jwt_openai = _make_jwt({"https://api.openai.com/profile": {"email": "custom@openai.com"}, "sub": "auth0|openai-456"})

    email1, sub1 = ProfileAuthManager.extract_jwt_identity(jwt_standard)
    email2, sub2 = ProfileAuthManager.extract_jwt_identity(jwt_openai)

    assert email1 == "standard@gmail.com"
    assert sub1 == "google-sub-123"
    assert email2 == "custom@openai.com"
    assert sub2 == "auth0|openai-456"





# ==============================================================================
# TEST I: OpenCode Go API key paste with whitespace / newline stripping
# ==============================================================================
def test_i_opencode_go_paste_sanitization():
    """Verify paste sanitization strips leading/trailing newlines and whitespace."""
    raw_key_with_whitespace = "  \n\t opencode-sk-test-secret-999 \r\n  "
    cleaned = raw_key_with_whitespace.strip()

    assert cleaned == "opencode-sk-test-secret-999"

    is_valid, masked, models = ProfileAuthManager.verify_opencode_token(cleaned)
    assert is_valid is True
    assert masked == "opencode-...-999"
    assert "opencode-go-3" in models


# ==============================================================================
# TEST J: Password masking does not corrupt backing value
# ==============================================================================
def test_j_password_masking_integrity():
    """Verify password masking (show='*') retains exact backing value."""
    secret = "sk-proj-actualsecretvalue123456"
    auth_data = {
        "provider": "openai-codex",
        "profile_id": "codex-worker-2",
        "auth_mode": "api_key",
        "api_key": secret,
    }
    ProfileAuthManager.save_profile_auth("openai-codex", "codex-worker-2", auth_data)

    loaded = ProfileAuthManager.load_profile_auth("openai-codex", "codex-worker-2")
    assert loaded["api_key"] == secret


# ==============================================================================
# TEST K: Zero-secret logging across Codex OAuth and API key flows
# ==============================================================================
def test_k_zero_secret_logging_codex(caplog):
    """Verify raw tokens and API keys are never written to logger."""
    import logging
    caplog.set_level(logging.DEBUG)

    session = CodexOAuthSession("codex-orch")
    session._finalize_with_tokens("super_secret_access_token_9999", "super_secret_refresh_token_8888")

    all_logs = " ".join([r.message for r in caplog.records])
    assert "super_secret_access_token_9999" not in all_logs
    assert "super_secret_refresh_token_8888" not in all_logs


# ==============================================================================
# TEST L: Single completion lock prevents double token exchange
# ==============================================================================
def test_l_single_completion_protection():
    """Verify second finalize call does not re-save or double process."""
    session = CodexOAuthSession("codex-orch")
    res1 = session._finalize_with_tokens("token1", "refresh1")
    assert res1 is True
    assert session._is_completed is True

    # Try manual input after completed
    ok, msg = session.handle_manual_input("some_other_token")
    assert ok is True
    assert "Авторизация уже успешно завершена" in msg


# ==============================================================================
# TEST M: Session cleanup on cancellation
# ==============================================================================
def test_m_codex_session_cleanup():
    """Verify cancel_codex_oauth_session cancels polling and cleans global store."""
    session_id, url, code = start_codex_oauth("codex-spare-1")
    session = get_codex_oauth_session(session_id)
    assert session is not None

    cancel_codex_oauth_session(session_id)
    assert session.status == "cancelled"
    assert get_codex_oauth_session(session_id) is None