pupilometer/tests/test_vision.py

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": "yolov10.onnx"}

    @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)
        mock_backend_class.assert_called_once_with(self.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)
        mock_backend_class.assert_called_once_with(self.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)
        mock_backend_class.assert_called_once_with(self.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()

    @patch('vision.logging')
    @patch.dict('sys.modules', {'onnxruntime': MagicMock(), 'onnxruntime-gpu': None})
    def test_python_backend_cpu_fallback(self, mock_logging):
        """
        Test that PythonBackend falls back to CPU when onnxruntime-gpu is not available.
        """
        mock_ort = sys.modules['onnxruntime']
        mock_ort.get_available_providers.return_value = ['CPUExecutionProvider']
        
        backend = PythonBackend(self.config)
        
        mock_logging.warning.assert_called_with("onnxruntime-gpu is not available or CUDA is not configured. Falling back to onnxruntime (CPU).")
        self.assertEqual(backend.ort, mock_ort)

    @patch('vision.logging')
    @patch.dict('sys.modules', {'onnxruntime': MagicMock()})
    def test_python_backend_gpu_selection(self, mock_logging):
        """
        Test that PythonBackend selects GPU when onnxruntime-gpu is available.
        """
        mock_ort_gpu = MagicMock()
        mock_ort_gpu.get_available_providers.return_value = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        sys.modules['onnxruntime'] = mock_ort_gpu
        
        backend = PythonBackend(self.config)
        
        mock_logging.info.assert_any_call("CUDA is available. Using onnxruntime-gpu.")
        self.assertEqual(backend.ort, mock_ort_gpu)

    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)

if __name__ == '__main__':
    unittest.main()