add basic tests

Author: deruyter92

pytest.ini

@@ -1,3 +1,9 @@
 [pytest]
 markers =
     functional: functional tests
+
+filterwarnings =
+    # Suppress NumPy deprecation warning from Keras/TensorFlow about np.object
+    ignore::FutureWarning:keras.*
+    ignore::FutureWarning:tensorflow.*
+    ignore:In the future `np.object` will be defined:FutureWarning

tests/test_config.py

@@ -0,0 +1,50 @@
+"""
+Tests for configuration file reading
+"""
+import pytest
+from pathlib import Path
+from dlclive.core import config
+
+
+class TestConfig:
+    """Test configuration file reading"""
+
+    def test_read_yaml_success(self, tmp_path):
+        """Test successfully reading a YAML config file"""
+        config_file = tmp_path / "pose_cfg.yaml"
+        # Write YAML content directly
+        yaml_content = """num_joints: 17
+all_joints:
+  - head
+  - neck
+  - shoulder
+batch_size: 1
+"""
+        config_file.write_text(yaml_content)
+        
+        result = config.read_yaml(config_file)
+        
+        assert result["num_joints"] == 17
+        assert result["all_joints"] == ["head", "neck", "shoulder"]
+        assert result["batch_size"] == 1
+
+    def test_read_yaml_nonexistent(self, tmp_path):
+        """Test that nonexistent config files raise FileNotFoundError"""
+        nonexistent = tmp_path / "nonexistent.yaml"
+        
+        with pytest.raises(FileNotFoundError):
+            config.read_yaml(nonexistent)
+
+    def test_read_yaml_path_resolution(self, tmp_path):
+        """Test that paths are properly resolved"""
+        config_file = tmp_path / "config.yaml"
+        config_file.write_text("test: value")
+        
+        # Test with relative path
+        result = config.read_yaml(str(config_file))
+        assert result["test"] == "value"
+        
+        # Test with Path object
+        result = config.read_yaml(config_file)
+        assert result["test"] == "value"
+

tests/test_dlclive.py

@@ -0,0 +1,207 @@
+"""
+Tests for DLCLive core functionality - frame processing, cropping, etc.
+"""
+import pytest
+import numpy as np
+from pathlib import Path
+from unittest.mock import Mock, MagicMock, patch
+from dlclive import DLCLive
+from dlclive.exceptions import DLCLiveError
+
+
+class TestDLCLive:
+    """Test DLCLive class core functionality"""
+
+    @pytest.fixture
+    def mock_runner(self):
+        """Create a mock runner for testing"""
+        runner = Mock()
+        runner.cfg = {"test": "config"}
+        runner.precision = "FP32"
+        runner.init_inference.return_value = np.zeros((17, 3))
+        runner.get_pose.return_value = np.zeros((17, 3))
+        runner.close.return_value = None
+        runner.read_config.return_value = {"test": "config"}
+        return runner
+
+    @pytest.fixture
+    def sample_frame(self):
+        """Create a sample frame for testing"""
+        return np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
+
+    @patch('dlclive.factory.build_runner')
+    def test_dlclive_initialization(self, mock_build_runner, mock_runner, tmp_path):
+        """Test DLCLive initialization"""
+        model_path = tmp_path / "model.pt"
+        model_path.write_text("test")
+        mock_build_runner.return_value = mock_runner
+        
+        dlc = DLCLive(model_path, model_type="pytorch")
+        
+        assert dlc.path == model_path
+        assert dlc.model_type == "pytorch"
+        assert dlc.is_initialized == False
+        assert dlc.cropping is None
+        assert dlc.dynamic == (False, 0.5, 10)
+        assert dlc.processor is None
+
+    @patch('dlclive.factory.build_runner')
+    def test_dlclive_cfg_property(self, mock_build_runner, mock_runner, tmp_path):
+        """Test accessing cfg property"""
+        model_path = tmp_path / "model.pt"
+        model_path.write_text("test")
+        mock_build_runner.return_value = mock_runner
+        
+        dlc = DLCLive(model_path)
+        assert dlc.cfg == {"test": "config"}
+
+    @patch('dlclive.factory.build_runner')
+    def test_dlclive_precision_property(self, mock_build_runner, mock_runner, tmp_path):
+        """Test accessing precision property"""
+        model_path = tmp_path / "model.pt"
+        model_path.write_text("test")
+        mock_build_runner.return_value = mock_runner
+        
+        dlc = DLCLive(model_path)
+        assert dlc.precision == "FP32"
+
+    @patch('dlclive.factory.build_runner')
+    def test_dlclive_read_config(self, mock_build_runner, mock_runner, tmp_path):
+        """Test reading configuration"""
+        model_path = tmp_path / "model.pt"
+        model_path.write_text("test")
+        mock_build_runner.return_value = mock_runner
+        
+        dlc = DLCLive(model_path)
+        config = dlc.read_config()
+        assert config == {"test": "config"}
+
+    @patch('dlclive.factory.build_runner')
+    def test_dlclive_parameterization(self, mock_build_runner, mock_runner, tmp_path):
+        """Test parameterization property"""
+        model_path = tmp_path / "model.pt"
+        model_path.write_text("test")
+        mock_build_runner.return_value = mock_runner
+        
+        dlc = DLCLive(model_path, cropping=[10, 100, 20, 200])
+        params = dlc.parameterization
+        
+        assert "path" in params
+        assert "cfg" in params
+        assert "model_type" in params
+        assert params["cropping"] == [10, 100, 20, 200]
+
+    @patch('dlclive.factory.build_runner')
+    @patch('dlclive.utils.img_to_rgb')
+    def test_process_frame_cropping(self, mock_img_to_rgb, mock_build_runner, 
+                                     mock_runner, sample_frame, tmp_path):
+        """Test frame processing with cropping"""
+        model_path = tmp_path / "model.pt"
+        model_path.write_text("test")
+        mock_build_runner.return_value = mock_runner
+        mock_img_to_rgb.side_effect = lambda x: x
+        
+        dlc = DLCLive(model_path, cropping=[10, 100, 20, 200])
+        result = dlc.process_frame(sample_frame)
+        
+        # Check that cropping was applied (result should be smaller)
+        assert result.shape[0] == 180  # 200 - 20
+        assert result.shape[1] == 90   # 100 - 10
+
+    @patch('dlclive.factory.build_runner')
+    @patch('dlclive.utils.resize_frame')
+    @patch('dlclive.utils.img_to_rgb')
+    def test_process_frame_resize(self, mock_img_to_rgb, mock_resize, 
+                                   mock_build_runner, mock_runner, sample_frame, tmp_path):
+        """Test frame processing with resize"""
+        model_path = tmp_path / "model.pt"
+        model_path.write_text("test")
+        mock_build_runner.return_value = mock_runner
+        mock_img_to_rgb.side_effect = lambda x: x
+        mock_resize.side_effect = lambda x, resize: x  # No actual resize in test
+        
+        dlc = DLCLive(model_path, resize=0.5)
+        dlc.process_frame(sample_frame)
+        
+        mock_resize.assert_called_once()
+
+    @patch('dlclive.factory.build_runner')
+    def test_process_frame_dynamic_cropping(self, mock_build_runner, mock_runner, 
+                                            sample_frame, tmp_path):
+        """Test dynamic cropping functionality"""
+        model_path = tmp_path / "model.pt"
+        model_path.write_text("test")
+        mock_build_runner.return_value = mock_runner
+        
+        # Create pose with detected body parts
+        pose = np.array([[100, 150, 0.8], [120, 160, 0.9], [80, 140, 0.7]])
+        
+        dlc = DLCLive(model_path, dynamic=(True, 0.5, 10))
+        dlc.pose = pose
+        
+        result = dlc.process_frame(sample_frame)
+        
+        # Check that dynamic cropping was applied
+        assert dlc.dynamic_cropping is not None
+        assert len(dlc.dynamic_cropping) == 4
+
+    @patch('dlclive.factory.build_runner')
+    def test_init_inference_no_frame(self, mock_build_runner, mock_runner, tmp_path):
+        """Test that init_inference raises error with no frame"""
+        model_path = tmp_path / "model.pt"
+        model_path.write_text("test")
+        mock_build_runner.return_value = mock_runner
+        
+        dlc = DLCLive(model_path)
+        
+        with pytest.raises(DLCLiveError, match="No frame provided"):
+            dlc.init_inference()
+
+    @patch('dlclive.factory.build_runner')
+    def test_get_pose_no_frame(self, mock_build_runner, mock_runner, tmp_path):
+        """Test that get_pose raises error with no frame"""
+        model_path = tmp_path / "model.pt"
+        model_path.write_text("test")
+        mock_build_runner.return_value = mock_runner
+        
+        dlc = DLCLive(model_path)
+        
+        with pytest.raises(DLCLiveError, match="No frame provided"):
+            dlc.get_pose()
+
+    @patch('dlclive.factory.build_runner')
+    def test_close(self, mock_build_runner, mock_runner, tmp_path):
+        """Test closing DLCLive instance"""
+        model_path = tmp_path / "model.pt"
+        model_path.write_text("test")
+        mock_build_runner.return_value = mock_runner
+        
+        dlc = DLCLive(model_path)
+        dlc.is_initialized = True
+        dlc.close()
+        
+        assert dlc.is_initialized == False
+        mock_runner.close.assert_called_once()
+
+    @patch('dlclive.factory.build_runner')
+    def test_post_process_pose_with_processor(self, mock_build_runner, mock_runner, 
+                                               sample_frame, tmp_path):
+        """Test pose post-processing with processor"""
+        model_path = tmp_path / "model.pt"
+        model_path.write_text("test")
+        mock_build_runner.return_value = mock_runner
+        
+        # Create mock processor
+        mock_processor = Mock()
+        mock_processor.process.return_value = np.ones((17, 3))
+        
+        dlc = DLCLive(model_path, processor=mock_processor)
+        dlc.pose = np.zeros((17, 3))
+        
+        # Manually call _post_process_pose
+        result = dlc._post_process_pose(sample_frame)
+        
+        mock_processor.process.assert_called_once()
+        np.testing.assert_array_equal(result, np.ones((17, 3)))
+
+

tests/test_engine.py

@@ -0,0 +1,66 @@
+"""
+Tests for the Engine class - engine detection and model type handling
+"""
+import pytest
+from pathlib import Path
+from dlclive.engine import Engine
+
+
+class TestEngine:
+    """Test Engine enum and detection methods"""
+
+    def test_engine_from_model_type_pytorch(self):
+        """Test detecting PyTorch engine from model type"""
+        assert Engine.from_model_type("pytorch") == Engine.PYTORCH
+        assert Engine.from_model_type("PyTorch") == Engine.PYTORCH
+        assert Engine.from_model_type("PYTORCH") == Engine.PYTORCH
+
+    def test_engine_from_model_type_tensorflow(self):
+        """Test detecting TensorFlow engine from model type"""
+        assert Engine.from_model_type("tensorflow") == Engine.TENSORFLOW
+        assert Engine.from_model_type("base") == Engine.TENSORFLOW
+        assert Engine.from_model_type("tensorrt") == Engine.TENSORFLOW
+        assert Engine.from_model_type("lite") == Engine.TENSORFLOW
+
+    def test_engine_from_model_type_invalid(self):
+        """Test that invalid model types raise ValueError"""
+        with pytest.raises(ValueError, match="Unknown model type"):
+            Engine.from_model_type("invalid")
+
+    def test_engine_from_model_path_tensorflow_dir(self, tmp_path):
+        """Test detecting TensorFlow engine from directory with .pb and pose_cfg.yaml"""
+        model_dir = tmp_path / "tensorflow_model"
+        model_dir.mkdir()
+        (model_dir / "pose_cfg.yaml").write_text("test")
+        (model_dir / "snapshot-100.pb").write_text("test")
+        
+        assert Engine.from_model_path(model_dir) == Engine.TENSORFLOW
+
+    def test_engine_from_model_path_pytorch_file(self, tmp_path):
+        """Test detecting PyTorch engine from .pt file"""
+        model_file = tmp_path / "model.pt"
+        model_file.write_text("test")
+        
+        assert Engine.from_model_path(model_file) == Engine.PYTORCH
+
+    def test_engine_from_model_path_nonexistent(self, tmp_path):
+        """Test that nonexistent paths raise FileNotFoundError"""
+        nonexistent = tmp_path / "nonexistent"
+        with pytest.raises(FileNotFoundError):
+            Engine.from_model_path(nonexistent)
+
+    def test_engine_from_model_path_invalid(self, tmp_path):
+        """Test that invalid model paths raise ValueError"""
+        # Directory without required files
+        invalid_dir = tmp_path / "invalid"
+        invalid_dir.mkdir()
+        with pytest.raises(ValueError, match="Could not determine engine"):
+            Engine.from_model_path(invalid_dir)
+        
+        # File with wrong extension
+        wrong_ext = tmp_path / "model.txt"
+        wrong_ext.write_text("test")
+        with pytest.raises(ValueError, match="Could not determine engine"):
+            Engine.from_model_path(wrong_ext)
+
+

tests/test_exceptions.py

@@ -0,0 +1,37 @@
+"""
+Tests for exception classes
+"""
+import pytest
+from dlclive.exceptions import DLCLiveError, DLCLiveWarning
+
+
+class TestExceptions:
+    """Test exception and warning classes"""
+
+    def test_dlclive_error(self):
+        """Test DLCLiveError can be raised and caught"""
+        with pytest.raises(DLCLiveError):
+            raise DLCLiveError("Test error message")
+
+    def test_dlclive_error_message(self):
+        """Test DLCLiveError preserves error message"""
+        error_msg = "Custom error message"
+        with pytest.raises(DLCLiveError, match=error_msg):
+            raise DLCLiveError(error_msg)
+
+    def test_dlclive_error_inheritance(self):
+        """Test DLCLiveError is an Exception"""
+        error = DLCLiveError("test")
+        assert isinstance(error, Exception)
+
+    def test_dlclive_warning(self):
+        """Test DLCLiveWarning can be issued"""
+        with pytest.warns(DLCLiveWarning):
+            import warnings
+            warnings.warn("Test warning", DLCLiveWarning)
+
+    def test_dlclive_warning_inheritance(self):
+        """Test DLCLiveWarning is a Warning"""
+        assert issubclass(DLCLiveWarning, Warning)
+
+