PerceptionMetrics/tests/test_model_loading.py at afb54aeb9eefa930fe35a016ba2af03e1c5fae2a · JdeRobot/PerceptionMetrics · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
"""Unit tests for model loading robustness"""
import os
import tempfile
import pytest

torch = pytest.importorskip("torch", reason="torch not installed in this environment")


from perceptionmetrics.models.torch_detection import TorchImageDetectionModel
from perceptionmetrics.models.torch_segmentation import TorchImageSegmentationModel


class TestModelLoadingExceptions:
    """Test that model loading raises appropriate exceptions"""

    @pytest.fixture
    def temp_files(self):
        """Create temporary test files"""
        with tempfile.TemporaryDirectory() as tmpdir:
            # Create dummy model file (corrupted)
            bad_model = os.path.join(tmpdir, "bad_model.pt")
            with open(bad_model, "w") as f:
                f.write("this is not a pytorch model")

            # Create dummy ontology
            ontology = os.path.join(tmpdir, "ontology.json")
            import json
            with open(ontology, "w") as f:
                json.dump({
                    "car": {"idx": 0, "rgb": [0, 0, 0]},
                    "person": {"idx": 1, "rgb": [255, 0, 0]},
                }, f)

            # Create dummy config
            config = os.path.join(tmpdir, "config.json")
            with open(config, "w") as f:
                json.dump({
                    "resize": {"width": 512, "height": 512},
                    "normalization": {
                        "mean": [0.485, 0.456, 0.406],
                        "std": [0.229, 0.224, 0.225]
                    },
                    "batch_size": 1,
                    "model_format": "torchvision"
                }, f)

            yield {
                "bad_model": bad_model,
                "ontology": ontology,
                "config": config,
                "tmpdir": tmpdir
            }

    def test_detection_model_bad_file_raises_specific_error(self, temp_files):
        """Test that loading corrupted model raises RuntimeError, not generic Exception"""
        with pytest.raises(RuntimeError) as exc_info:
            TorchImageDetectionModel(
                model=temp_files["bad_model"],
                model_cfg=temp_files["config"],
                ontology_fname=temp_files["ontology"]
            )

        # Check error message is informative
        error_msg = str(exc_info.value)
        assert "Failed to load model" in error_msg
        assert "TorchScript error" in error_msg or "PyTorch error" in error_msg

    def test_segmentation_model_bad_file_raises_specific_error(self, temp_files):
        """Test that loading corrupted segmentation model raises RuntimeError"""
        with pytest.raises(RuntimeError) as exc_info:
            TorchImageSegmentationModel(
                model=temp_files["bad_model"],
                model_cfg=temp_files["config"],
                ontology_fname=temp_files["ontology"]
            )

        error_msg = str(exc_info.value)
        assert "Failed to load model" in error_msg

    def test_detection_model_missing_file_raises_file_not_found(self, temp_files):
        """Test that missing model file raises FileNotFoundError"""
        with pytest.raises(FileNotFoundError) as exc_info:
            TorchImageDetectionModel(
                model="/nonexistent/path/model.pt",
                model_cfg=temp_files["config"],
                ontology_fname=temp_files["ontology"]
            )

        assert "Model file not found" in str(exc_info.value)

    def test_segmentation_model_missing_file_raises_file_not_found(self, temp_files):
        """Test that missing segmentation model file raises FileNotFoundError"""
        with pytest.raises(FileNotFoundError) as exc_info:
            TorchImageSegmentationModel(
                model="/nonexistent/path/model.pt",
                model_cfg=temp_files["config"],
                ontology_fname=temp_files["ontology"]
            )

        assert "Model file not found" in str(exc_info.value)