From 7665af94dca65a7f39ce095eab92ea4b762e386b Mon Sep 17 00:00:00 2001
From: spencerwooo <spencer.woo@outlook.com>
Date: Mon, 15 Apr 2024 22:02:44 +0800
Subject: [PATCH] Add unit tests

---
 .vscode/settings.json             |  7 ++++++
 pyproject.toml                    |  2 +-
 test/test_feature_layer.py        | 14 +++++++++++
 test/test_feature_layers.py       | 33 ++++++++++++++++++++++++
 test/test_get_layer_candidates.py | 42 +++++++++++++++++++++++++++++++
 5 files changed, 97 insertions(+), 1 deletion(-)
 create mode 100644 .vscode/settings.json
 create mode 100644 test/test_feature_layer.py
 create mode 100644 test/test_feature_layers.py
 create mode 100644 test/test_get_layer_candidates.py

diff --git a/.vscode/settings.json b/.vscode/settings.json
new file mode 100644
index 0000000..3e99ede
--- /dev/null
+++ b/.vscode/settings.json
@@ -0,0 +1,7 @@
+{
+    "python.testing.pytestArgs": [
+        "."
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
+}
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 1b2c741..5f648fb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,7 +27,7 @@ Documentation = "https://github.com/spencerwooo/torch-featurelayer/blob/main/REA
 
 [project.optional-dependencies]
 dev = ["mypy", "ruff"]
-test = ["torchvision>=0.15,<0.18"]
+test = ["torchvision>=0.15,<0.18", "pytest"]
 
 [build-system]
 requires = ["setuptools"]
diff --git a/test/test_feature_layer.py b/test/test_feature_layer.py
new file mode 100644
index 0000000..c93ab2f
--- /dev/null
+++ b/test/test_feature_layer.py
@@ -0,0 +1,14 @@
+import torch
+from torchvision.models import alexnet
+from torch_featurelayer import FeatureLayer
+
+
+def test_feature_layer():
+    model = alexnet()
+    hooked_model = FeatureLayer(model, 'features.12')
+
+    x = torch.randn(1, 3, 224, 224)
+    feature_output, output = hooked_model(x)
+
+    assert feature_output.shape == (1, 256, 6, 6)
+    assert output.shape == (1, 1000)
diff --git a/test/test_feature_layers.py b/test/test_feature_layers.py
new file mode 100644
index 0000000..c2982aa
--- /dev/null
+++ b/test/test_feature_layers.py
@@ -0,0 +1,33 @@
+import torch
+from torchvision.models import alexnet
+from torch_featurelayer import FeatureLayers
+
+
+def test_feature_layers():
+    model = alexnet()
+    layer_paths = [
+        'features.2',
+        'features.5',
+        'features.9',
+        'features.12',
+        'avgpool',
+        'classifier.2',
+        'classifier.4',
+    ]
+    hooked_model = FeatureLayers(model, layer_paths)
+
+    x = torch.randn(1, 3, 224, 224)
+    feature_outputs, output = hooked_model(x)
+
+    feature_output_shapes = {
+        'features.2': (1, 64, 27, 27),
+        'features.5': (1, 192, 13, 13),
+        'features.9': (1, 256, 13, 13),
+        'features.12': (1, 256, 6, 6),
+        'avgpool': (1, 256, 6, 6),
+        'classifier.2': (1, 4096),
+        'classifier.4': (1, 4096),
+    }
+    for layer_path, feature_output in feature_outputs.items():
+        assert feature_output.shape == feature_output_shapes[layer_path]
+    assert output.shape == (1, 1000)
diff --git a/test/test_get_layer_candidates.py b/test/test_get_layer_candidates.py
new file mode 100644
index 0000000..29cdbd4
--- /dev/null
+++ b/test/test_get_layer_candidates.py
@@ -0,0 +1,42 @@
+from torchvision.models import alexnet
+from torch_featurelayer import get_layer_candidates
+
+
+def test_get_layer_candidates_depth1():
+    model = alexnet()
+    candidates = list(get_layer_candidates(model, max_depth=1))
+    assert candidates == [
+        'features',
+        'avgpool',
+        'classifier',
+    ]
+
+
+def test_get_layer_candidates_depth2():
+    model = alexnet()
+    candidates = list(get_layer_candidates(model, max_depth=2))
+    assert candidates == [
+        'features',
+        'features.0',
+        'features.1',
+        'features.2',
+        'features.3',
+        'features.4',
+        'features.5',
+        'features.6',
+        'features.7',
+        'features.8',
+        'features.9',
+        'features.10',
+        'features.11',
+        'features.12',
+        'avgpool',
+        'classifier',
+        'classifier.0',
+        'classifier.1',
+        'classifier.2',
+        'classifier.3',
+        'classifier.4',
+        'classifier.5',
+        'classifier.6',
+    ]