45 lines
1.8 KiB
Python
45 lines
1.8 KiB
Python
"""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"
|