tracer-cloud--opensre/tests/interactive_shell/test_bedrock_model.py

"""Tests for Bedrock-specific model validation and custom ID handling."""

from __future__ import annotations

from dataclasses import dataclass

import pytest

from surfaces.interactive_shell.command_registry.model.command import (
    _prompt_custom_model_id,
    _reasoning_model_menu_choices,
    _toolcall_model_menu_choices,
)
from surfaces.interactive_shell.command_registry.model.switching import _is_model_allowed


@dataclass(frozen=True)
class _FakeModelOption:
    value: str
    label: str = ""


@dataclass(frozen=True)
class _FakeProvider:
    value: str
    models: tuple[_FakeModelOption, ...] = ()
    allow_custom_models: bool = False


# ─────────────────────────────────────────────────────────────────────────────
# _is_model_allowed — custom model bypass
# ─────────────────────────────────────────────────────────────────────────────


class TestIsModelSupportedBedrock:
    """Bedrock must accept any non-empty model string (ARNs, regional prefixes, etc.)."""

    def test_bedrock_accepts_inference_profile_id(self) -> None:
        provider = _FakeProvider(value="bedrock", allow_custom_models=True)
        assert _is_model_allowed(provider, "us.anthropic.claude-sonnet-4-6") is True

    def test_bedrock_accepts_eu_prefix(self) -> None:
        provider = _FakeProvider(value="bedrock", allow_custom_models=True)
        assert _is_model_allowed(provider, "eu.anthropic.claude-sonnet-4-6") is True

    def test_bedrock_accepts_global_prefix(self) -> None:
        provider = _FakeProvider(value="bedrock", allow_custom_models=True)
        assert _is_model_allowed(provider, "global.anthropic.claude-opus-4-7") is True

    def test_bedrock_accepts_on_demand_model(self) -> None:
        provider = _FakeProvider(value="bedrock", allow_custom_models=True)
        assert _is_model_allowed(provider, "mistral.mistral-large-3-675b-instruct") is True

    def test_bedrock_accepts_full_arn(self) -> None:
        arn = "arn:aws:bedrock:us-east-1:123456789012:inference-profile/my-profile"
        provider = _FakeProvider(value="bedrock", allow_custom_models=True)
        assert _is_model_allowed(provider, arn) is True

    def test_bedrock_rejects_empty_string(self) -> None:
        provider = _FakeProvider(value="bedrock", allow_custom_models=True)
        assert _is_model_allowed(provider, "") is False

    def test_ollama_also_accepts_any_model(self) -> None:
        """Ollama shares the same bypass pattern — verify it still works."""
        provider = _FakeProvider(value="ollama", allow_custom_models=True)
        assert _is_model_allowed(provider, "llama3.2") is True

    def test_other_provider_requires_match(self) -> None:
        """Non-bypass providers must match the curated model list."""
        models = (_FakeModelOption(value="gpt-5.4"),)
        provider = _FakeProvider(value="anthropic", models=models)
        assert _is_model_allowed(provider, "gpt-5.4") is True
        assert _is_model_allowed(provider, "gpt-unknown") is False

    def test_custom_provider_accepts_unknown_non_empty_model(self) -> None:
        provider = _FakeProvider(value="openai", allow_custom_models=True)
        assert _is_model_allowed(provider, "gpt-5.5") is True


# ─────────────────────────────────────────────────────────────────────────────
# Menu choices — __custom__ sentinel
# ─────────────────────────────────────────────────────────────────────────────


class TestMenuChoicesCustomOption:
    """Bedrock menus must include the __custom__ escape hatch."""

    def test_reasoning_menu_includes_custom_for_bedrock(self) -> None:
        provider = _FakeProvider(
            value="bedrock",
            models=(_FakeModelOption(value="us.anthropic.claude-sonnet-4-6"),),
            allow_custom_models=True,
        )
        choices = _reasoning_model_menu_choices(provider)
        values = [v for v, _ in choices]
        assert "__custom__" in values

    def test_toolcall_menu_includes_custom_for_bedrock(self) -> None:
        provider = _FakeProvider(
            value="bedrock",
            models=(_FakeModelOption(value="us.anthropic.claude-sonnet-4-6"),),
            allow_custom_models=True,
        )
        choices = _toolcall_model_menu_choices(provider)
        values = [v for v, _ in choices]
        assert "__custom__" in values

    def test_reasoning_menu_includes_custom_for_openai(self) -> None:
        provider = _FakeProvider(
            value="openai",
            models=(_FakeModelOption(value="gpt-5.4"),),
            allow_custom_models=True,
        )
        choices = _reasoning_model_menu_choices(provider)
        values = [v for v, _ in choices]
        assert "__custom__" in values

    def test_toolcall_menu_includes_custom_for_openai(self) -> None:
        provider = _FakeProvider(
            value="openai",
            models=(_FakeModelOption(value="gpt-5.4"),),
            allow_custom_models=True,
        )
        choices = _toolcall_model_menu_choices(provider)
        values = [v for v, _ in choices]
        assert "__custom__" in values


# ─────────────────────────────────────────────────────────────────────────────
# _prompt_custom_model_id
# ─────────────────────────────────────────────────────────────────────────────


class TestPromptCustomModelId:
    """Custom model ID prompt edge cases."""

    def test_returns_stripped_value(self, monkeypatch: pytest.MonkeyPatch) -> None:
        from rich.console import Console

        console = Console(force_terminal=False)
        monkeypatch.setattr(console, "input", lambda _prompt: "  us.anthropic.claude-opus-4-7  ")
        result = _prompt_custom_model_id(console)
        assert result == "us.anthropic.claude-opus-4-7"

    def test_returns_none_on_empty_input(self, monkeypatch: pytest.MonkeyPatch) -> None:
        from rich.console import Console

        console = Console(force_terminal=False)
        monkeypatch.setattr(console, "input", lambda _prompt: "   ")
        result = _prompt_custom_model_id(console)
        assert result is None

    def test_returns_none_on_keyboard_interrupt(self, monkeypatch: pytest.MonkeyPatch) -> None:
        from rich.console import Console

        console = Console(force_terminal=False)

        def _raise_interrupt(_prompt: str) -> str:
            raise KeyboardInterrupt

        monkeypatch.setattr(console, "input", _raise_interrupt)
        result = _prompt_custom_model_id(console)
        assert result is None

    def test_returns_none_on_eof(self, monkeypatch: pytest.MonkeyPatch) -> None:
        from rich.console import Console

        console = Console(force_terminal=False)

        def _raise_eof(_prompt: str) -> str:
            raise EOFError

        monkeypatch.setattr(console, "input", _raise_eof)
        result = _prompt_custom_model_id(console)
        assert result is None


# ─────────────────────────────────────────────────────────────────────────────
# Bedrock ProviderOption in wizard config
# ─────────────────────────────────────────────────────────────────────────────


class TestBedrockProviderConfig:
    """Verify Bedrock is registered correctly in the wizard config."""

    def test_bedrock_in_supported_providers(self) -> None:
        from surfaces.cli.wizard.config import PROVIDER_BY_VALUE

        assert "bedrock" in PROVIDER_BY_VALUE

    def test_bedrock_credential_kind_is_none(self) -> None:
        from surfaces.cli.wizard.config import PROVIDER_BY_VALUE

        provider = PROVIDER_BY_VALUE["bedrock"]
        assert provider.credential_kind == "none"

    def test_bedrock_has_curated_models(self) -> None:
        from surfaces.cli.wizard.config import PROVIDER_BY_VALUE

        provider = PROVIDER_BY_VALUE["bedrock"]
        assert len(provider.models) >= 10

    def test_bedrock_curated_models_use_inference_profiles(self) -> None:
        """All Claude models in the curated list must use us.* inference profile IDs."""
        from surfaces.cli.wizard.config import PROVIDER_BY_VALUE

        provider = PROVIDER_BY_VALUE["bedrock"]
        claude_models = [m for m in provider.models if "anthropic" in str(getattr(m, "value", ""))]
        for model in claude_models:
            value = str(getattr(model, "value", ""))
            assert value.startswith("us."), (
                f"Claude model '{value}' must use a us.* inference profile ID"
            )

    def test_bedrock_has_toolcall_model_env(self) -> None:
        from surfaces.cli.wizard.config import PROVIDER_BY_VALUE

        provider = PROVIDER_BY_VALUE["bedrock"]
        assert provider.toolcall_model_env == "BEDROCK_TOOLCALL_MODEL"

    def test_bedrock_api_key_env_is_empty(self) -> None:
        """api_key_env="" is intentional — Bedrock uses IAM auth, not an API key."""
        from surfaces.cli.wizard.config import PROVIDER_BY_VALUE

        provider = PROVIDER_BY_VALUE["bedrock"]
        assert provider.api_key_env == ""
        # Empty string must be falsy so downstream ``bool(provider.api_key_env)``
        # checks correctly skip API-key validation for Bedrock.
        assert not provider.api_key_env

    def test_bedrock_no_credential_default(self) -> None:
        """credential_default must use the dataclass default (empty string).

        Region is picked up from AWS_DEFAULT_REGION / ~/.aws/config, not from
        the wizard credential prompt (which is skipped for credential_kind="none").
        """
        from surfaces.cli.wizard.config import PROVIDER_BY_VALUE

        provider = PROVIDER_BY_VALUE["bedrock"]
        assert provider.credential_default == ""


# ─────────────────────────────────────────────────────────────────────────────
# _interactive_set_toolcall + __custom__ integration
# ─────────────────────────────────────────────────────────────────────────────


class TestInteractiveSetToolcallCustom:
    """Verify the __custom__ branch inside _interactive_set_toolcall wires through correctly."""

    def test_custom_toolcall_calls_switch_with_typed_id(
        self, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """Selecting __custom__ in toolcall menu should prompt and call switch_toolcall_model."""
        from unittest.mock import patch

        from rich.console import Console

        from surfaces.interactive_shell.command_registry.model import command as model_mod

        console = Console(force_terminal=False)
        custom_id = "eu.anthropic.claude-sonnet-4-6"

        # First call: pick provider "bedrock"; second call: pick "__custom__"
        choose_returns = iter(["bedrock", "__custom__"])
        monkeypatch.setattr(model_mod, "repl_choose_one", lambda **_kw: next(choose_returns))
        monkeypatch.setattr(model_mod, "_prompt_custom_model_id", lambda *_args: custom_id)

        with patch.object(model_mod, "switch_toolcall_model", return_value=True) as mock_switch:
            result = model_mod._interactive_set_toolcall(console)

        assert result is True
        mock_switch.assert_called_once_with(custom_id, console, provider_name="bedrock")

    def test_custom_toolcall_returns_none_on_cancel(self, monkeypatch: pytest.MonkeyPatch) -> None:
        """If user cancels the custom prompt, _interactive_set_toolcall returns None."""
        from rich.console import Console

        from surfaces.interactive_shell.command_registry.model import command as model_mod

        console = Console(force_terminal=False)

        choose_returns = iter(["bedrock", "__custom__"])
        monkeypatch.setattr(model_mod, "repl_choose_one", lambda **_kw: next(choose_returns))
        monkeypatch.setattr(model_mod, "_prompt_custom_model_id", lambda *_args: None)

        result = model_mod._interactive_set_toolcall(console)
        assert result is None