File size: 2,188 Bytes
d49de5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
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']
|