"""Tests for main module functions."""

import pytest
from src.main import format_model_name


def test_format_model_name_with_provider_prefix():
    """Test formatting model names with provider prefix."""
    assert format_model_name("anthropic/claude-3-haiku") == "Claude-3-Haiku"
    assert format_model_name("meta-llama/llama-3.1-8b-instruct") == "Llama-3.1-8B-Instruct"
    assert format_model_name("google/gemini-pro") == "Gemini-Pro"


def test_format_model_name_without_prefix():
    """Test formatting model names without provider prefix."""
    assert format_model_name("llama-2-7b") == "Llama-2-7B"
    assert format_model_name("gpt-4-turbo") == "Gpt-4-Turbo"
    assert format_model_name("mistral-7b-instruct") == "Mistral-7B-Instruct"


def test_format_model_name_preserves_version_numbers():
    """Test that version numbers are preserved correctly."""
    assert format_model_name("llama-3.1-8b") == "Llama-3.1-8B"
    assert format_model_name("claude-3-opus") == "Claude-3-Opus"
    assert format_model_name("gpt-4") == "Gpt-4"


def test_format_model_name_handles_underscores():
    """Test that underscores are handled like hyphens."""
    assert format_model_name("model_name_test") == "Model-Name-Test"
    assert format_model_name("llama_2_7b") == "Llama-2-7B"


def test_format_model_name_uppercase_size_designators():
    """Test that size designators ending in 'b' are uppercased."""
    assert format_model_name("model-8b") == "Model-8B"
    assert format_model_name("model-70b") == "Model-70B"
    assert format_model_name("model-13b-instruct") == "Model-13B-Instruct"


def test_format_model_name_single_word():
    """Test formatting single-word model names."""
    assert format_model_name("gemini") == "Gemini"
    assert format_model_name("claude") == "Claude"