conftest.py

"""
Pytest configuration and fixtures for Lambda Bedrock tests
"""

import json
import pytest
from unittest.mock import Mock, patch
import os
import sys

# Add the project root to the Python path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))


@pytest.fixture
def mock_context():
    """Create a mock Lambda context for testing"""
    context = Mock()
    context.aws_request_id = "test-request-id-12345"
    context.function_name = "test-bedrock-function"
    context.function_version = "1"
    context.invoked_function_arn = "arn:aws:lambda:us-east-1:123456789012:function:test-bedrock-function"
    context.memory_limit_in_mb = 128
    context.remaining_time_in_millis = 30000
    context.log_group_name = "/aws/lambda/test-bedrock-function"
    context.log_stream_name = "2023/01/01/[$LATEST]test-stream"
    return context


@pytest.fixture
def sample_chat_event():
    """Create a sample chat completion event"""
    return {
        'body': json.dumps({
            'messages': [
                {'role': 'user', 'content': 'Hello, how are you?'}
            ],
            'model': 'custom-model',
            'max_tokens': 100,
            'temperature': 0.7,
            'top_p': 0.9
        }),
        'headers': {'Content-Type': 'application/json'}
    }


@pytest.fixture
def conversation_event():
    """Create a conversation with history event"""
    return {
        'body': json.dumps({
            'messages': [
                {'role': 'user', 'content': 'My name is John'},
                {'role': 'assistant', 'content': 'Hello John! Nice to meet you.'},
                {'role': 'user', 'content': "What's my name?"}
            ],
            'model': 'custom-model'
        }),
        'headers': {'Content-Type': 'application/json'}
    }


@pytest.fixture
def image_chat_event():
    """Create a chat event with image content"""
    return {
        'body': json.dumps({
            'messages': [
                {
                    'role': 'user',
                    'content': [
                        {'type': 'text', 'text': 'What do you see in this image?'},
                        {
                            'type': 'image_url',
                            'image_url': {
                                'url': 'data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAYEBQYFBAYGBQYHBwYIChAKCgkJChQODwwQFxQYGBcUFhYaHSUfGhsjHBYWICwgIyYnKSopGR8tMC0oMCUoKSj/2wBDAQcHBwoIChMKChMoGhYaKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCgoKCj/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAv/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwCdABmX/9k='
                            }
                        }
                    ]
                }
            ],
            'model': 'custom-model'
        }),
        'headers': {'Content-Type': 'application/json'}
    }


@pytest.fixture
def error_event():
    """Create an event that should cause an error"""
    return {
        'body': json.dumps({
            'model': 'custom-model'
            # Missing messages array
        }),
        'headers': {'Content-Type': 'application/json'}
    }


@pytest.fixture
def mock_bedrock_response():
    """Create a mock Bedrock response"""
    response = Mock()
    response['body'].read.return_value = json.dumps({
        'generation': 'Hello! I am doing well, thank you for asking. How can I help you today?'
    }).encode('utf-8')
    return response


@pytest.fixture
def mock_bedrock_client(mock_bedrock_response):
    """Create a mock Bedrock client"""
    with patch('lambda_function.boto3.client') as mock_client:
        mock_bedrock = Mock()
        mock_bedrock.invoke_model.return_value = mock_bedrock_response
        mock_client.return_value = mock_bedrock
        yield mock_bedrock


@pytest.fixture
def test_events():
    """Load test events from JSON file"""
    try:
        with open('test_events.json', 'r') as f:
            return json.load(f)
    except FileNotFoundError:
        pytest.skip("test_events.json not found")


# Pytest configuration
def pytest_configure(config):
    """Configure pytest with custom markers"""
    config.addinivalue_line(
        "markers", "unit: mark test as a unit test"
    )
    config.addinivalue_line(
        "markers", "integration: mark test as an integration test"
    )
    config.addinivalue_line(
        "markers", "bedrock: mark test as requiring Bedrock mocking"
    )
    config.addinivalue_line(
        "markers", "slow: mark test as slow running"
    )


def pytest_collection_modifyitems(config, items):
    """Modify test collection to add markers automatically"""
    for item in items:
        # Add unit marker to test files that start with test_
        if item.nodeid.startswith("tests/test_"):
            item.add_marker(pytest.mark.unit)
        
        # Add bedrock marker to tests that use Bedrock mocking
        if "bedrock" in item.nodeid.lower() or "mock_bedrock" in str(item.function):
            item.add_marker(pytest.mark.bedrock)