tests/test_pipeline.py

import pytest
from unittest.mock import patch, MagicMock
from found_protocol.pipeline import FoundProtocolPipeline

@pytest.fixture
def mock_genai_models():
    with patch('google.generativeai.GenerativeModel') as mock_model_class:
        mock_vision = MagicMock()
        mock_vision.generate_content.return_value.text = "Mock vision description."
        
        mock_f1 = MagicMock()
        mock_f1.generate_content.return_value.text = '{"cli_output": "[MOCK F1] Analysis.", "detected_concepts": [], "narrative_impact": "Mock."}'
        
        mock_f2 = MagicMock()
        mock_f2.generate_content.return_value.text = '{"cli_output": "[MOCK F2] Compilation.", "emotional_state": "mock", "system_action": "Mock."}'

        def mock_side_effect(model_name, **kwargs):
            if 'vision' in model_name: return mock_vision
            if 'FOUND-1' in kwargs.get('system_instruction', ''): return mock_f1
            if 'FOUND-2' in kwargs.get('system_instruction', ''): return mock_f2
            return MagicMock()

        mock_model_class.side_effect = mock_side_effect
        yield mock_model_class

def test_pipeline_initialization(mock_genai_models):
    """Tests if the pipeline initializes correctly."""
    pipeline = FoundProtocolPipeline(gemini_api_key="TEST_KEY")
    assert pipeline is not None
    assert "INIT" in pipeline.state.memory_log[0]

def test_pipeline_run_mocked(tmp_path, mock_genai_models):
    """Tests a full run with mocked API calls."""
    dummy_video = tmp_path / "dummy.mp4"
    dummy_video.write_bytes(b'dummy data')

    with patch('cv2.VideoCapture') as mock_cap_class:
        mock_cap = MagicMock()
        mock_cap.get.return_value = 100
        mock_cap.read.return_value = (True, "fake_frame")
        mock_cap_class.return_value = mock_cap
        
        pipeline = FoundProtocolPipeline(gemini_api_key="TEST_KEY")
        result = pipeline.run(str(dummy_video))
        
        assert "error" not in result
        assert "found1_output" in result
        assert isinstance(result['found1_output'], dict)
        assert "cli_output" in result['found1_output']
        assert "[MOCK F1]" in result['found1_output']['cli_output']