921 lines
No EOL
32 KiB
Python
921 lines
No EOL
32 KiB
Python
"""
|
|
Unit tests for pipeline execution functionality.
|
|
"""
|
|
import pytest
|
|
import asyncio
|
|
import numpy as np
|
|
from unittest.mock import Mock, MagicMock, patch, AsyncMock
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import json
|
|
|
|
from detector_worker.pipeline.pipeline_executor import (
|
|
PipelineExecutor,
|
|
PipelineContext,
|
|
PipelineResult,
|
|
BranchResult,
|
|
ExecutionMode
|
|
)
|
|
from detector_worker.detection.detection_result import DetectionResult, BoundingBox
|
|
from detector_worker.core.exceptions import PipelineError, ModelError, ActionError
|
|
|
|
|
|
class TestPipelineContext:
|
|
"""Test pipeline context data structure."""
|
|
|
|
def test_creation(self):
|
|
"""Test pipeline context creation."""
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=np.zeros((480, 640, 3), dtype=np.uint8)
|
|
)
|
|
|
|
assert context.camera_id == "camera_001"
|
|
assert context.display_id == "display_001"
|
|
assert context.session_id == "session_123"
|
|
assert context.timestamp == 1640995200000
|
|
assert context.frame_data.shape == (480, 640, 3)
|
|
assert context.metadata == {}
|
|
assert context.crop_region is None
|
|
|
|
def test_creation_with_crop_region(self):
|
|
"""Test context creation with crop region."""
|
|
crop_region = (100, 200, 300, 400)
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=np.zeros((480, 640, 3), dtype=np.uint8),
|
|
crop_region=crop_region
|
|
)
|
|
|
|
assert context.crop_region == crop_region
|
|
|
|
def test_add_metadata(self):
|
|
"""Test adding metadata to context."""
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=np.zeros((480, 640, 3), dtype=np.uint8)
|
|
)
|
|
|
|
context.add_metadata("model_id", "yolo_v8")
|
|
context.add_metadata("confidence_threshold", 0.8)
|
|
|
|
assert context.metadata["model_id"] == "yolo_v8"
|
|
assert context.metadata["confidence_threshold"] == 0.8
|
|
|
|
def test_get_cropped_frame(self):
|
|
"""Test getting cropped frame."""
|
|
frame = np.ones((480, 640, 3), dtype=np.uint8) * 255
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=frame,
|
|
crop_region=(100, 200, 300, 400)
|
|
)
|
|
|
|
cropped = context.get_cropped_frame()
|
|
|
|
assert cropped.shape == (200, 200, 3) # 400-200, 300-100
|
|
assert np.all(cropped == 255)
|
|
|
|
def test_get_cropped_frame_no_crop(self):
|
|
"""Test getting frame when no crop region specified."""
|
|
frame = np.ones((480, 640, 3), dtype=np.uint8) * 255
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=frame
|
|
)
|
|
|
|
cropped = context.get_cropped_frame()
|
|
|
|
assert np.array_equal(cropped, frame)
|
|
|
|
|
|
class TestBranchResult:
|
|
"""Test branch execution result."""
|
|
|
|
def test_creation_success(self):
|
|
"""Test successful branch result creation."""
|
|
detections = [
|
|
DetectionResult("car", 0.9, BoundingBox(100, 200, 300, 400), 1001, 1640995200000)
|
|
]
|
|
|
|
result = BranchResult(
|
|
branch_id="car_brand_cls",
|
|
success=True,
|
|
detections=detections,
|
|
metadata={"brand": "Toyota"},
|
|
execution_time=0.15
|
|
)
|
|
|
|
assert result.branch_id == "car_brand_cls"
|
|
assert result.success is True
|
|
assert len(result.detections) == 1
|
|
assert result.metadata["brand"] == "Toyota"
|
|
assert result.execution_time == 0.15
|
|
assert result.error is None
|
|
|
|
def test_creation_failure(self):
|
|
"""Test failed branch result creation."""
|
|
result = BranchResult(
|
|
branch_id="car_brand_cls",
|
|
success=False,
|
|
error="Model inference failed",
|
|
execution_time=0.05
|
|
)
|
|
|
|
assert result.branch_id == "car_brand_cls"
|
|
assert result.success is False
|
|
assert result.detections == []
|
|
assert result.metadata == {}
|
|
assert result.error == "Model inference failed"
|
|
|
|
|
|
class TestPipelineResult:
|
|
"""Test pipeline execution result."""
|
|
|
|
def test_creation_success(self):
|
|
"""Test successful pipeline result creation."""
|
|
main_detections = [
|
|
DetectionResult("car", 0.9, BoundingBox(100, 200, 300, 400), 1001, 1640995200000)
|
|
]
|
|
|
|
branch_results = {
|
|
"car_brand_cls": BranchResult("car_brand_cls", True, [], {"brand": "Toyota"}, 0.1),
|
|
"car_bodytype_cls": BranchResult("car_bodytype_cls", True, [], {"body_type": "Sedan"}, 0.12)
|
|
}
|
|
|
|
result = PipelineResult(
|
|
success=True,
|
|
detections=main_detections,
|
|
branch_results=branch_results,
|
|
total_execution_time=0.5
|
|
)
|
|
|
|
assert result.success is True
|
|
assert len(result.detections) == 1
|
|
assert len(result.branch_results) == 2
|
|
assert result.total_execution_time == 0.5
|
|
assert result.error is None
|
|
|
|
def test_creation_failure(self):
|
|
"""Test failed pipeline result creation."""
|
|
result = PipelineResult(
|
|
success=False,
|
|
error="Pipeline execution failed",
|
|
total_execution_time=0.1
|
|
)
|
|
|
|
assert result.success is False
|
|
assert result.detections == []
|
|
assert result.branch_results == {}
|
|
assert result.error == "Pipeline execution failed"
|
|
|
|
def test_get_combined_results(self):
|
|
"""Test getting combined results from all branches."""
|
|
main_detections = [
|
|
DetectionResult("car", 0.9, BoundingBox(100, 200, 300, 400), 1001, 1640995200000)
|
|
]
|
|
|
|
branch_results = {
|
|
"car_brand_cls": BranchResult("car_brand_cls", True, [], {"brand": "Toyota"}, 0.1),
|
|
"car_bodytype_cls": BranchResult("car_bodytype_cls", True, [], {"body_type": "Sedan"}, 0.12)
|
|
}
|
|
|
|
result = PipelineResult(
|
|
success=True,
|
|
detections=main_detections,
|
|
branch_results=branch_results,
|
|
total_execution_time=0.5
|
|
)
|
|
|
|
combined = result.get_combined_results()
|
|
|
|
assert "brand" in combined
|
|
assert "body_type" in combined
|
|
assert combined["brand"] == "Toyota"
|
|
assert combined["body_type"] == "Sedan"
|
|
|
|
|
|
class TestPipelineExecutor:
|
|
"""Test pipeline execution functionality."""
|
|
|
|
def test_initialization(self):
|
|
"""Test pipeline executor initialization."""
|
|
executor = PipelineExecutor()
|
|
|
|
assert isinstance(executor.thread_pool, ThreadPoolExecutor)
|
|
assert executor.max_workers == 4
|
|
assert executor.execution_mode == ExecutionMode.PARALLEL
|
|
assert executor.timeout == 30.0
|
|
|
|
def test_initialization_custom_config(self):
|
|
"""Test initialization with custom configuration."""
|
|
config = {
|
|
"max_workers": 8,
|
|
"execution_mode": "sequential",
|
|
"timeout": 60.0
|
|
}
|
|
|
|
executor = PipelineExecutor(config)
|
|
|
|
assert executor.max_workers == 8
|
|
assert executor.execution_mode == ExecutionMode.SEQUENTIAL
|
|
assert executor.timeout == 60.0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_pipeline_simple(self, mock_yolo_model, mock_frame):
|
|
"""Test simple pipeline execution."""
|
|
# Mock pipeline configuration
|
|
pipeline_config = {
|
|
"modelId": "car_detection_v1",
|
|
"modelFile": "car_detection.pt",
|
|
"expectedClasses": ["car"],
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.8,
|
|
"branches": [],
|
|
"actions": []
|
|
}
|
|
|
|
# Mock detection result
|
|
mock_result = Mock()
|
|
mock_result.boxes = Mock()
|
|
mock_result.boxes.data = torch.tensor([
|
|
[100, 200, 300, 400, 0.9, 0]
|
|
])
|
|
mock_result.boxes.id = torch.tensor([1001])
|
|
|
|
mock_yolo_model.track.return_value = [mock_result]
|
|
|
|
executor = PipelineExecutor()
|
|
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=mock_frame
|
|
)
|
|
|
|
with patch('detector_worker.models.model_manager.ModelManager') as mock_model_manager:
|
|
mock_model_manager.return_value.get_model.return_value = mock_yolo_model
|
|
|
|
result = await executor.execute_pipeline(pipeline_config, context)
|
|
|
|
assert result.success is True
|
|
assert len(result.detections) == 1
|
|
assert result.detections[0].class_name == "0" # Default class name
|
|
assert result.detections[0].confidence == 0.9
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_pipeline_with_branches(self, mock_yolo_model, mock_frame):
|
|
"""Test pipeline execution with classification branches."""
|
|
import torch
|
|
|
|
# Mock main detection
|
|
mock_detection_result = Mock()
|
|
mock_detection_result.boxes = Mock()
|
|
mock_detection_result.boxes.data = torch.tensor([
|
|
[100, 200, 300, 400, 0.9, 0] # car detection
|
|
])
|
|
mock_detection_result.boxes.id = torch.tensor([1001])
|
|
|
|
# Mock classification results
|
|
mock_brand_result = Mock()
|
|
mock_brand_result.probs = Mock()
|
|
mock_brand_result.probs.top1 = 2 # Toyota
|
|
mock_brand_result.probs.top1conf = 0.85
|
|
|
|
mock_bodytype_result = Mock()
|
|
mock_bodytype_result.probs = Mock()
|
|
mock_bodytype_result.probs.top1 = 1 # Sedan
|
|
mock_bodytype_result.probs.top1conf = 0.78
|
|
|
|
mock_yolo_model.track.return_value = [mock_detection_result]
|
|
mock_yolo_model.predict.return_value = [mock_brand_result]
|
|
|
|
mock_brand_model = Mock()
|
|
mock_brand_model.predict.return_value = [mock_brand_result]
|
|
mock_brand_model.names = {0: "Honda", 1: "Ford", 2: "Toyota"}
|
|
|
|
mock_bodytype_model = Mock()
|
|
mock_bodytype_model.predict.return_value = [mock_bodytype_result]
|
|
mock_bodytype_model.names = {0: "SUV", 1: "Sedan", 2: "Hatchback"}
|
|
|
|
# Pipeline configuration with branches
|
|
pipeline_config = {
|
|
"modelId": "car_detection_v1",
|
|
"modelFile": "car_detection.pt",
|
|
"expectedClasses": ["car"],
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.8,
|
|
"branches": [
|
|
{
|
|
"modelId": "car_brand_cls",
|
|
"modelFile": "car_brand.pt",
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.7,
|
|
"parallel": True,
|
|
"crop": True,
|
|
"cropClass": "car"
|
|
},
|
|
{
|
|
"modelId": "car_bodytype_cls",
|
|
"modelFile": "car_bodytype.pt",
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.7,
|
|
"parallel": True,
|
|
"crop": True,
|
|
"cropClass": "car"
|
|
}
|
|
],
|
|
"actions": []
|
|
}
|
|
|
|
executor = PipelineExecutor()
|
|
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=mock_frame
|
|
)
|
|
|
|
with patch('detector_worker.models.model_manager.ModelManager') as mock_model_manager:
|
|
def get_model_side_effect(model_id, camera_id):
|
|
if model_id == "car_detection_v1":
|
|
return mock_yolo_model
|
|
elif model_id == "car_brand_cls":
|
|
return mock_brand_model
|
|
elif model_id == "car_bodytype_cls":
|
|
return mock_bodytype_model
|
|
return None
|
|
|
|
mock_model_manager.return_value.get_model.side_effect = get_model_side_effect
|
|
|
|
result = await executor.execute_pipeline(pipeline_config, context)
|
|
|
|
assert result.success is True
|
|
assert len(result.detections) == 1
|
|
assert len(result.branch_results) == 2
|
|
|
|
# Check branch results
|
|
assert "car_brand_cls" in result.branch_results
|
|
assert "car_bodytype_cls" in result.branch_results
|
|
|
|
brand_result = result.branch_results["car_brand_cls"]
|
|
assert brand_result.success is True
|
|
assert brand_result.metadata.get("brand") == "Toyota"
|
|
|
|
bodytype_result = result.branch_results["car_bodytype_cls"]
|
|
assert bodytype_result.success is True
|
|
assert bodytype_result.metadata.get("body_type") == "Sedan"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_pipeline_sequential_mode(self, mock_yolo_model, mock_frame):
|
|
"""Test pipeline execution in sequential mode."""
|
|
import torch
|
|
|
|
config = {"execution_mode": "sequential"}
|
|
executor = PipelineExecutor(config)
|
|
|
|
# Mock detection result
|
|
mock_result = Mock()
|
|
mock_result.boxes = Mock()
|
|
mock_result.boxes.data = torch.tensor([
|
|
[100, 200, 300, 400, 0.9, 0]
|
|
])
|
|
mock_result.boxes.id = torch.tensor([1001])
|
|
|
|
mock_yolo_model.track.return_value = [mock_result]
|
|
|
|
pipeline_config = {
|
|
"modelId": "car_detection_v1",
|
|
"modelFile": "car_detection.pt",
|
|
"expectedClasses": ["car"],
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.8,
|
|
"branches": [
|
|
{
|
|
"modelId": "car_brand_cls",
|
|
"modelFile": "car_brand.pt",
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.7,
|
|
"parallel": False # Sequential execution
|
|
}
|
|
],
|
|
"actions": []
|
|
}
|
|
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=mock_frame
|
|
)
|
|
|
|
with patch('detector_worker.models.model_manager.ModelManager') as mock_model_manager:
|
|
mock_model_manager.return_value.get_model.return_value = mock_yolo_model
|
|
|
|
result = await executor.execute_pipeline(pipeline_config, context)
|
|
|
|
assert result.success is True
|
|
assert executor.execution_mode == ExecutionMode.SEQUENTIAL
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_pipeline_with_actions(self, mock_yolo_model, mock_frame):
|
|
"""Test pipeline execution with actions."""
|
|
import torch
|
|
|
|
# Mock detection result
|
|
mock_result = Mock()
|
|
mock_result.boxes = Mock()
|
|
mock_result.boxes.data = torch.tensor([
|
|
[100, 200, 300, 400, 0.9, 0]
|
|
])
|
|
mock_result.boxes.id = torch.tensor([1001])
|
|
|
|
mock_yolo_model.track.return_value = [mock_result]
|
|
|
|
# Pipeline configuration with actions
|
|
pipeline_config = {
|
|
"modelId": "car_detection_v1",
|
|
"modelFile": "car_detection.pt",
|
|
"expectedClasses": ["car"],
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.8,
|
|
"branches": [],
|
|
"actions": [
|
|
{
|
|
"type": "redis_save_image",
|
|
"region": "car",
|
|
"key": "inference:{display_id}:{timestamp}:{session_id}",
|
|
"expire_seconds": 600
|
|
},
|
|
{
|
|
"type": "postgresql_insert",
|
|
"table": "detections",
|
|
"fields": {
|
|
"camera_id": "{camera_id}",
|
|
"detection_class": "{class}",
|
|
"confidence": "{confidence}"
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
executor = PipelineExecutor()
|
|
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=mock_frame
|
|
)
|
|
|
|
with patch('detector_worker.models.model_manager.ModelManager') as mock_model_manager, \
|
|
patch('detector_worker.pipeline.action_executor.ActionExecutor') as mock_action_executor:
|
|
|
|
mock_model_manager.return_value.get_model.return_value = mock_yolo_model
|
|
mock_action_executor.return_value.execute_actions = AsyncMock(return_value=True)
|
|
|
|
result = await executor.execute_pipeline(pipeline_config, context)
|
|
|
|
assert result.success is True
|
|
# Actions should be executed
|
|
mock_action_executor.return_value.execute_actions.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_pipeline_model_error(self, mock_frame):
|
|
"""Test pipeline execution with model error."""
|
|
pipeline_config = {
|
|
"modelId": "car_detection_v1",
|
|
"modelFile": "car_detection.pt",
|
|
"expectedClasses": ["car"],
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.8,
|
|
"branches": [],
|
|
"actions": []
|
|
}
|
|
|
|
executor = PipelineExecutor()
|
|
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=mock_frame
|
|
)
|
|
|
|
with patch('detector_worker.models.model_manager.ModelManager') as mock_model_manager:
|
|
# Model manager raises error
|
|
mock_model_manager.return_value.get_model.side_effect = ModelError("Model not found")
|
|
|
|
result = await executor.execute_pipeline(pipeline_config, context)
|
|
|
|
assert result.success is False
|
|
assert "Model not found" in result.error
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_pipeline_timeout(self, mock_yolo_model, mock_frame):
|
|
"""Test pipeline execution timeout."""
|
|
import torch
|
|
|
|
# Configure short timeout
|
|
config = {"timeout": 0.001} # Very short timeout
|
|
executor = PipelineExecutor(config)
|
|
|
|
# Mock slow model inference
|
|
def slow_inference(*args, **kwargs):
|
|
import time
|
|
time.sleep(1) # Longer than timeout
|
|
mock_result = Mock()
|
|
mock_result.boxes = None
|
|
return [mock_result]
|
|
|
|
mock_yolo_model.track.side_effect = slow_inference
|
|
|
|
pipeline_config = {
|
|
"modelId": "car_detection_v1",
|
|
"modelFile": "car_detection.pt",
|
|
"expectedClasses": ["car"],
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.8,
|
|
"branches": [],
|
|
"actions": []
|
|
}
|
|
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=mock_frame
|
|
)
|
|
|
|
with patch('detector_worker.models.model_manager.ModelManager') as mock_model_manager:
|
|
mock_model_manager.return_value.get_model.return_value = mock_yolo_model
|
|
|
|
result = await executor.execute_pipeline(pipeline_config, context)
|
|
|
|
assert result.success is False
|
|
assert "timeout" in result.error.lower()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_branch_parallel(self, mock_frame):
|
|
"""Test parallel branch execution."""
|
|
import torch
|
|
|
|
# Mock classification model
|
|
mock_brand_model = Mock()
|
|
mock_result = Mock()
|
|
mock_result.probs = Mock()
|
|
mock_result.probs.top1 = 1
|
|
mock_result.probs.top1conf = 0.85
|
|
mock_brand_model.predict.return_value = [mock_result]
|
|
mock_brand_model.names = {0: "Honda", 1: "Toyota", 2: "Ford"}
|
|
|
|
executor = PipelineExecutor()
|
|
|
|
# Branch configuration
|
|
branch_config = {
|
|
"modelId": "car_brand_cls",
|
|
"modelFile": "car_brand.pt",
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.7,
|
|
"parallel": True,
|
|
"crop": True,
|
|
"cropClass": "car"
|
|
}
|
|
|
|
# Mock detected regions
|
|
regions = {
|
|
"car": {
|
|
"bbox": [100, 200, 300, 400],
|
|
"confidence": 0.9,
|
|
"detection": DetectionResult("car", 0.9, BoundingBox(100, 200, 300, 400), 1001)
|
|
}
|
|
}
|
|
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=mock_frame
|
|
)
|
|
|
|
with patch('detector_worker.models.model_manager.ModelManager') as mock_model_manager:
|
|
mock_model_manager.return_value.get_model.return_value = mock_brand_model
|
|
|
|
result = await executor._execute_branch(branch_config, regions, context)
|
|
|
|
assert result.success is True
|
|
assert result.branch_id == "car_brand_cls"
|
|
assert result.metadata.get("brand") == "Toyota"
|
|
assert result.execution_time > 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_branch_no_trigger_class(self, mock_frame):
|
|
"""Test branch execution when trigger class not detected."""
|
|
executor = PipelineExecutor()
|
|
|
|
branch_config = {
|
|
"modelId": "car_brand_cls",
|
|
"modelFile": "car_brand.pt",
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.7
|
|
}
|
|
|
|
# No car detected
|
|
regions = {
|
|
"truck": {
|
|
"bbox": [100, 200, 300, 400],
|
|
"confidence": 0.9,
|
|
"detection": DetectionResult("truck", 0.9, BoundingBox(100, 200, 300, 400), 1002)
|
|
}
|
|
}
|
|
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=mock_frame
|
|
)
|
|
|
|
result = await executor._execute_branch(branch_config, regions, context)
|
|
|
|
assert result.success is False
|
|
assert "trigger class not detected" in result.error.lower()
|
|
|
|
def test_wait_for_branches(self):
|
|
"""Test waiting for specific branches to complete."""
|
|
executor = PipelineExecutor()
|
|
|
|
# Mock completed branch results
|
|
branch_results = {
|
|
"car_brand_cls": BranchResult("car_brand_cls", True, [], {"brand": "Toyota"}, 0.1),
|
|
"car_bodytype_cls": BranchResult("car_bodytype_cls", True, [], {"body_type": "Sedan"}, 0.12),
|
|
"license_ocr": BranchResult("license_ocr", True, [], {"license": "ABC123"}, 0.2)
|
|
}
|
|
|
|
# Wait for specific branches
|
|
wait_for = ["car_brand_cls", "car_bodytype_cls"]
|
|
completed = executor._wait_for_branches(branch_results, wait_for, timeout=1.0)
|
|
|
|
assert completed is True
|
|
|
|
# Wait for non-existent branch (should timeout)
|
|
wait_for_missing = ["car_brand_cls", "nonexistent_branch"]
|
|
completed = executor._wait_for_branches(branch_results, wait_for_missing, timeout=0.1)
|
|
|
|
assert completed is False
|
|
|
|
def test_validate_pipeline_config(self):
|
|
"""Test pipeline configuration validation."""
|
|
executor = PipelineExecutor()
|
|
|
|
# Valid configuration
|
|
valid_config = {
|
|
"modelId": "car_detection_v1",
|
|
"modelFile": "car_detection.pt",
|
|
"expectedClasses": ["car"],
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.8
|
|
}
|
|
|
|
assert executor._validate_pipeline_config(valid_config) is True
|
|
|
|
# Invalid configuration (missing required fields)
|
|
invalid_config = {
|
|
"modelFile": "car_detection.pt"
|
|
# Missing modelId
|
|
}
|
|
|
|
with pytest.raises(PipelineError):
|
|
executor._validate_pipeline_config(invalid_config)
|
|
|
|
def test_crop_frame_for_detection(self, mock_frame):
|
|
"""Test frame cropping for detection."""
|
|
executor = PipelineExecutor()
|
|
|
|
detection = DetectionResult("car", 0.9, BoundingBox(100, 200, 300, 400), 1001)
|
|
|
|
cropped = executor._crop_frame_for_detection(mock_frame, detection)
|
|
|
|
assert cropped.shape == (200, 200, 3) # 400-200, 300-100
|
|
|
|
def test_crop_frame_invalid_bounds(self, mock_frame):
|
|
"""Test frame cropping with invalid bounds."""
|
|
executor = PipelineExecutor()
|
|
|
|
# Detection outside frame bounds
|
|
detection = DetectionResult("car", 0.9, BoundingBox(-100, -200, 50, 100), 1001)
|
|
|
|
cropped = executor._crop_frame_for_detection(mock_frame, detection)
|
|
|
|
# Should handle bounds gracefully
|
|
assert cropped.shape[0] > 0
|
|
assert cropped.shape[1] > 0
|
|
|
|
|
|
class TestPipelineExecutorPerformance:
|
|
"""Test pipeline executor performance and optimization."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parallel_branch_execution_performance(self, mock_frame):
|
|
"""Test that parallel execution is faster than sequential."""
|
|
import time
|
|
import torch
|
|
|
|
def slow_inference(*args, **kwargs):
|
|
time.sleep(0.1) # Simulate slow inference
|
|
mock_result = Mock()
|
|
mock_result.probs = Mock()
|
|
mock_result.probs.top1 = 1
|
|
mock_result.probs.top1conf = 0.85
|
|
return [mock_result]
|
|
|
|
mock_model = Mock()
|
|
mock_model.predict.side_effect = slow_inference
|
|
mock_model.names = {0: "Class0", 1: "Class1"}
|
|
|
|
# Test parallel execution
|
|
parallel_executor = PipelineExecutor({"execution_mode": "parallel", "max_workers": 2})
|
|
|
|
branch_configs = [
|
|
{
|
|
"modelId": f"branch_{i}",
|
|
"modelFile": f"branch_{i}.pt",
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.7,
|
|
"parallel": True
|
|
}
|
|
for i in range(3) # 3 branches
|
|
]
|
|
|
|
regions = {
|
|
"car": {
|
|
"bbox": [100, 200, 300, 400],
|
|
"confidence": 0.9,
|
|
"detection": DetectionResult("car", 0.9, BoundingBox(100, 200, 300, 400), 1001)
|
|
}
|
|
}
|
|
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=mock_frame
|
|
)
|
|
|
|
with patch('detector_worker.models.model_manager.ModelManager') as mock_model_manager:
|
|
mock_model_manager.return_value.get_model.return_value = mock_model
|
|
|
|
start_time = time.time()
|
|
|
|
# Execute branches in parallel
|
|
tasks = [
|
|
parallel_executor._execute_branch(config, regions, context)
|
|
for config in branch_configs
|
|
]
|
|
|
|
results = await asyncio.gather(*tasks)
|
|
|
|
parallel_time = time.time() - start_time
|
|
|
|
# Parallel execution should be faster than 3 * 0.1 seconds
|
|
assert parallel_time < 0.25 # Allow some overhead
|
|
assert len(results) == 3
|
|
assert all(result.success for result in results)
|
|
|
|
def test_thread_pool_management(self):
|
|
"""Test thread pool creation and management."""
|
|
# Test different worker counts
|
|
for workers in [1, 2, 4, 8]:
|
|
executor = PipelineExecutor({"max_workers": workers})
|
|
assert executor.max_workers == workers
|
|
assert executor.thread_pool._max_workers == workers
|
|
|
|
def test_memory_management_large_frames(self):
|
|
"""Test memory management with large frames."""
|
|
executor = PipelineExecutor()
|
|
|
|
# Create large frame
|
|
large_frame = np.ones((1080, 1920, 3), dtype=np.uint8) * 128
|
|
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=large_frame,
|
|
crop_region=(500, 400, 1000, 800)
|
|
)
|
|
|
|
# Get cropped frame
|
|
cropped = context.get_cropped_frame()
|
|
|
|
# Should reduce memory usage
|
|
assert cropped.shape == (400, 500, 3) # Much smaller than original
|
|
assert cropped.nbytes < large_frame.nbytes
|
|
|
|
|
|
class TestPipelineExecutorErrorHandling:
|
|
"""Test comprehensive error handling."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_branch_execution_error_isolation(self, mock_frame):
|
|
"""Test that errors in one branch don't affect others."""
|
|
executor = PipelineExecutor()
|
|
|
|
# Mock models - one fails, one succeeds
|
|
failing_model = Mock()
|
|
failing_model.predict.side_effect = Exception("Model crashed")
|
|
|
|
success_model = Mock()
|
|
mock_result = Mock()
|
|
mock_result.probs = Mock()
|
|
mock_result.probs.top1 = 1
|
|
mock_result.probs.top1conf = 0.85
|
|
success_model.predict.return_value = [mock_result]
|
|
success_model.names = {0: "Class0", 1: "Class1"}
|
|
|
|
branch_configs = [
|
|
{
|
|
"modelId": "failing_branch",
|
|
"modelFile": "failing.pt",
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.7,
|
|
"parallel": True
|
|
},
|
|
{
|
|
"modelId": "success_branch",
|
|
"modelFile": "success.pt",
|
|
"triggerClasses": ["car"],
|
|
"minConfidence": 0.7,
|
|
"parallel": True
|
|
}
|
|
]
|
|
|
|
regions = {
|
|
"car": {
|
|
"bbox": [100, 200, 300, 400],
|
|
"confidence": 0.9,
|
|
"detection": DetectionResult("car", 0.9, BoundingBox(100, 200, 300, 400), 1001)
|
|
}
|
|
}
|
|
|
|
context = PipelineContext(
|
|
camera_id="camera_001",
|
|
display_id="display_001",
|
|
session_id="session_123",
|
|
timestamp=1640995200000,
|
|
frame_data=mock_frame
|
|
)
|
|
|
|
def get_model_side_effect(model_id, camera_id):
|
|
if model_id == "failing_branch":
|
|
return failing_model
|
|
elif model_id == "success_branch":
|
|
return success_model
|
|
return None
|
|
|
|
with patch('detector_worker.models.model_manager.ModelManager') as mock_model_manager:
|
|
mock_model_manager.return_value.get_model.side_effect = get_model_side_effect
|
|
|
|
# Execute branches
|
|
tasks = [
|
|
executor._execute_branch(config, regions, context)
|
|
for config in branch_configs
|
|
]
|
|
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
# One should fail, one should succeed
|
|
failing_result = next(r for r in results if isinstance(r, BranchResult) and r.branch_id == "failing_branch")
|
|
success_result = next(r for r in results if isinstance(r, BranchResult) and r.branch_id == "success_branch")
|
|
|
|
assert failing_result.success is False
|
|
assert "Model crashed" in failing_result.error
|
|
|
|
assert success_result.success is True
|
|
assert success_result.error is None |