tests/test_model_adapters.py

from __future__ import annotations

from prompt_injection_framework.models.groq import _extract_retry_delay_seconds
from prompt_injection_framework.models.ollama import OllamaModelAdapter
from prompt_injection_framework.orchestration.config import ExperimentConfig
from prompt_injection_framework.orchestration.runner import create_model_adapter


def _config(provider: str, model_name: str) -> ExperimentConfig:
    from pathlib import Path

    return ExperimentConfig(
        name="adapter_test",
        description="",
        seed=42,
        sample_size=1,
        output_dir=Path("results"),
        dataset_path=Path("data/normalized/attack_cases.jsonl"),
        provider=provider,
        model_name=model_name,
        max_tokens=32,
        temperature=0.0,
        execution_mode="plain",
        mitigation_enabled=False,
        mitigation_strategy="none",
        mitigation_options={},
        include_sources=(),
        include_attack_surfaces=(),
        include_attack_families=(),
        include_attack_subtypes=(),
        require_expected_target=False,
    )


def test_create_model_adapter_supports_ollama() -> None:
    adapter = create_model_adapter(_config("ollama", "llama3.1:8b"))
    assert isinstance(adapter, OllamaModelAdapter)
    assert adapter.model_name == "llama3.1:8b"


def test_ollama_adapter_uses_default_local_base_url() -> None:
    adapter = OllamaModelAdapter(model_name="qwen2.5:7b")
    assert adapter.api_url == "http://127.0.0.1:11434/api/generate"


def test_groq_retry_delay_parses_seconds() -> None:
    delay = _extract_retry_delay_seconds("Please try again in 3.47s.", attempt=1)
    assert delay == 3.47


def test_groq_retry_delay_parses_milliseconds() -> None:
    delay = _extract_retry_delay_seconds("Please try again in 990ms.", attempt=1)
    assert delay == 0.99


def test_groq_retry_delay_falls_back_to_attempt_backoff() -> None:
    delay = _extract_retry_delay_seconds("rate_limit_exceeded", attempt=3)
    assert delay == 6.0