pupilometer/tests/test_vision.py
2025-11-28 16:18:30 +07:00

136 lines
5.5 KiB
Python

import unittest
from unittest.mock import patch, MagicMock
import sys
import os
import numpy as np
# Add the src/controllerSoftware directory to the Python path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../src/controllerSoftware')))
# Mock the gi module
sys.modules['gi'] = MagicMock()
sys.modules['gi.repository'] = MagicMock()
from vision import VisionSystem, DeepStreamBackend, PythonBackend, MockBackend
class TestVisionSystem(unittest.TestCase):
"""
Unit tests for the VisionSystem class.
"""
def setUp(self):
"""
Set up a VisionSystem instance with a mocked backend for each test.
"""
self.config = {"camera_id": 0, "model_path": "yolov8n-seg.pt"}
@patch('platform.system', return_value='Linux')
@patch('vision.DeepStreamBackend')
def test_initialization_linux(self, mock_backend_class, mock_system):
"""
Test that the VisionSystem initializes the DeepStreamBackend on Linux.
"""
mock_backend_instance = mock_backend_class.return_value
vision_system = VisionSystem(self.config)
expected_config = self.config.copy()
expected_config.setdefault('model_name', 'yolov8n-seg.pt') # Add default model_name
mock_backend_class.assert_called_once_with(expected_config)
self.assertEqual(vision_system._backend, mock_backend_instance)
@patch('platform.system', return_value='Windows')
@patch('vision.DeepStreamBackend')
def test_initialization_windows(self, mock_backend_class, mock_system):
"""
Test that the VisionSystem initializes the DeepStreamBackend on Windows.
"""
mock_backend_instance = mock_backend_class.return_value
vision_system = VisionSystem(self.config)
expected_config = self.config.copy()
expected_config.setdefault('model_name', 'yolov8n-seg.pt') # Add default model_name
mock_backend_class.assert_called_once_with(expected_config)
self.assertEqual(vision_system._backend, mock_backend_instance)
@patch('platform.system', return_value='Darwin')
@patch('vision.PythonBackend')
def test_initialization_macos(self, mock_backend_class, mock_system):
"""
Test that the VisionSystem initializes the PythonBackend on macOS.
"""
mock_backend_instance = mock_backend_class.return_value
vision_system = VisionSystem(self.config)
expected_config = self.config.copy()
expected_config.setdefault('model_name', 'yolov8n-seg.pt') # Add default model_name
mock_backend_class.assert_called_once_with(expected_config)
self.assertEqual(vision_system._backend, mock_backend_instance)
@patch('platform.system', return_value='UnsupportedOS')
def test_initialization_unsupported(self, mock_system):
"""
Test that the VisionSystem raises an exception on an unsupported OS.
"""
with self.assertRaises(NotImplementedError):
VisionSystem(self.config)
@patch('platform.system', return_value='Linux')
@patch('vision.DeepStreamBackend')
def test_start(self, mock_backend_class, mock_system):
"""
Test that the start method calls the backend's start method.
"""
mock_backend_instance = mock_backend_class.return_value
vision_system = VisionSystem(self.config)
vision_system.start()
mock_backend_instance.start.assert_called_once()
@patch('platform.system', return_value='Linux')
@patch('vision.DeepStreamBackend')
def test_stop(self, mock_backend_class, mock_system):
"""
Test that the stop method calls the backend's stop method.
"""
mock_backend_instance = mock_backend_class.return_value
vision_system = VisionSystem(self.config)
vision_system.stop()
mock_backend_instance.stop.assert_called_once()
@patch('platform.system', return_value='Linux')
@patch('vision.DeepStreamBackend')
def test_get_pupil_data(self, mock_backend_class, mock_system):
"""
Test that the get_pupil_data method calls the backend's get_pupil_data method.
"""
mock_backend_instance = mock_backend_class.return_value
vision_system = VisionSystem(self.config)
vision_system.get_pupil_data()
mock_backend_instance.get_pupil_data.assert_called_once()
@patch('platform.system', return_value='Linux')
@patch('vision.DeepStreamBackend')
def test_get_annotated_frame(self, mock_backend_class, mock_system):
"""
Test that the get_annotated_frame method calls the backend's get_annotated_frame method.
"""
mock_backend_instance = mock_backend_class.return_value
vision_system = VisionSystem(self.config)
vision_system.get_annotated_frame()
mock_backend_instance.get_annotated_frame.assert_called_once()
def test_mock_backend_methods(self):
"""
Test the methods of the MockBackend.
"""
backend = MockBackend(self.config)
backend.start()
backend.stop()
data = backend.get_pupil_data()
self.assertIn("pupil_position", data)
frame = backend.get_annotated_frame()
self.assertIsInstance(frame, np.ndarray)
def test_model_exists(self):
"""
Tests that the YOLO model file (.pt) exists at the expected location.
"""
model_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../src/controllerSoftware', self.config['model_path']))
self.assertTrue(os.path.exists(model_path), f"YOLO model file not found at {model_path}")