|
|
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'] |
|
|
|
