Add exception tests and assertions

.vscode/settings.json

@@ -1,6 +1,9 @@
 {
     "python.testing.pytestArgs": [
-        "."
+        "test",
+        "--color=yes",
+        "-o",
+        "log_cli=true"
     ],
     "python.testing.unittestEnabled": false,
     "python.testing.pytestEnabled": true

src/torch_featurelayer/feature_layer.py

@@ -21,6 +21,7 @@ class FeatureLayer:
     """
 
     def __init__(self, model: torch.nn.Module, feature_layer_path: str):
+        assert isinstance(model, torch.nn.Module), 'model must be a torch.nn.Module'
         self._model: torch.nn.Module = model
         self.feature_layer_path: str = feature_layer_path
         self.feature_layer_output: torch.Tensor | None = None  # output of the feature layer (must be global)

src/torch_featurelayer/feature_layers.py

@@ -19,6 +19,7 @@ class FeatureLayers:
     """
 
     def __init__(self, model: torch.nn.Module, feature_layer_paths: list[str]):
+        assert isinstance(model, torch.nn.Module), 'model must be a torch.nn.Module'
         self._model: torch.nn.Module = model
         self.feature_layer_paths: list[str] = feature_layer_paths
         self.feature_layer_outputs: dict[str, torch.Tensor | None] = dict().fromkeys(feature_layer_paths, None)

src/torch_featurelayer/layer_candidates.py

@@ -13,6 +13,9 @@ def get_layer_candidates(module: nn.Module, max_depth: int = 1) -> Generator[str
         Generator: A generator of the model's layer candidates in dot notation.
     """
 
+    assert max_depth >= 0, 'max_depth must be a non-negative integer'
+    assert isinstance(module, nn.Module), 'model must be a torch.nn.Module'
+
     def get_modules(model: nn.Module, prefix: str = '', depth: int = 0) -> Generator[str, None, None]:
         if prefix:
             yield prefix

test/test_feature_layer.py

@@ -1,6 +1,7 @@
+import pytest
 import torch
-from torchvision.models import alexnet
 from torch_featurelayer import FeatureLayer
+from torchvision.models import alexnet
 
 
 def test_feature_layer():
@@ -12,3 +13,22 @@ def test_feature_layer():
 
     assert feature_output.shape == (1, 256, 6, 6)
     assert output.shape == (1, 1000)
+
+
+def test_feature_layer_invalid_model():
+    with pytest.raises(AssertionError) as e:
+        model = 'invalid_model'
+        _ = FeatureLayer(model, 'features.12')
+
+    assert str(e.value) == 'model must be a torch.nn.Module'
+
+
+def test_feature_layer_nonexistent_layer_path():
+    with pytest.raises(AttributeError) as e:
+        model = alexnet()
+        hooked_model = FeatureLayer(model, 'path.to.blah')
+
+        x = torch.randn(1, 3, 224, 224)
+        _ = hooked_model(x)
+
+    assert str(e.value) == 'Layer path.to.blah not found in model.'

test/test_feature_layers.py

@@ -1,6 +1,7 @@
+import pytest
 import torch
-from torchvision.models import alexnet
 from torch_featurelayer import FeatureLayers
+from torchvision.models import alexnet
 
 
 def test_feature_layers():
@@ -31,3 +32,33 @@ def test_feature_layers():
     for layer_path, feature_output in feature_outputs.items():
         assert feature_output.shape == feature_output_shapes[layer_path]
     assert output.shape == (1, 1000)
+
+
+def test_feature_layers_invalid_model():
+    with pytest.raises(AssertionError) as e:
+        model = 'invalid_model'
+        _ = FeatureLayers(model, ['features.2', 'features.5', 'features.9'])
+    assert str(e.value) == 'model must be a torch.nn.Module'
+
+
+def test_feature_layers_contain_nonexistent_layer_path():
+    model = alexnet()
+    hooked_model = FeatureLayers(model, ['features.1', 'path.to.blah'])
+
+    x = torch.randn(1, 3, 224, 224)
+    feature_outputs, outputs = hooked_model(x)
+
+    assert feature_outputs['features.1'] is not None
+    assert feature_outputs['path.to.blah'] is None  # nonexistent layer path ignored
+    assert outputs is not None
+
+
+def test_feature_layers_empty_layer_paths():
+    model = alexnet()
+    hooked_model = FeatureLayers(model, [])
+
+    x = torch.randn(1, 3, 224, 224)
+    feature_outputs, output = hooked_model(x)
+
+    assert feature_outputs == {}
+    assert output is not None

test/test_get_layer_candidates.py

@@ -1,15 +1,12 @@
-from torchvision.models import alexnet
+import pytest
 from torch_featurelayer import get_layer_candidates
+from torchvision.models import alexnet
 
 
 def test_get_layer_candidates_depth1():
     model = alexnet()
     candidates = list(get_layer_candidates(model, max_depth=1))
-    assert candidates == [
-        'features',
-        'avgpool',
-        'classifier',
-    ]
+    assert candidates == ['features', 'avgpool', 'classifier']
 
 
 def test_get_layer_candidates_depth2():
@@ -40,3 +37,17 @@ def test_get_layer_candidates_depth2():
         'classifier.5',
         'classifier.6',
     ]
+
+
+def test_get_layer_candidates_invalid_model():
+    with pytest.raises(AssertionError) as e:
+        model = 'invalid_model'
+        get_layer_candidates(model, max_depth=1)
+    assert str(e.value) == 'model must be a torch.nn.Module'
+
+
+def test_get_layer_candidates_negative_depth():
+    with pytest.raises(AssertionError) as e:
+        model = alexnet()
+        get_layer_candidates(model, max_depth=-1)
+    assert str(e.value) == 'max_depth must be a non-negative integer'