compact_ai_model/tests/conftest.py

#!/usr/bin/env python3
"""
Pytest configuration and fixtures.
"""

import pytest
import sys
from pathlib import Path

# Add project root to path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))


@pytest.fixture
def sample_config():
    """Provide a sample configuration for testing."""
    from compact_ai_model.configs.config import Config
    return Config.get_tiny_config()


@pytest.fixture
def tiny_model(sample_config):
    """Provide a tiny model for testing."""
    from compact_ai_model.architecture.model import CompactAIModel
    model = CompactAIModel(sample_config.model, sample_config.thinking)
    model.eval()
    return model


@pytest.fixture
def sample_batch():
    """Provide a sample input batch for testing."""
    import torch
    batch_size, seq_len = 2, 16
    vocab_size = 1000  # Small vocab for testing
    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
    return input_ids