add a test

pyronear · Jul 31, 2023 · 42a2acc · 42a2acc
1 parent 1ad8b6c
commit 42a2acc
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tests/test_vision.py b/tests/test_vision.py
@@ -19,4 +19,8 @@ def test_classifier(mock_wildfire_image):
     # Test mask
     mask = np.ones((384, 640))
     out = model(mock_wildfire_image, mask)
+    assert out.shape == (1, 5)
+
+    mask = np.zeros((384, 640))
+    out = model(mock_wildfire_image, mask)
     assert out.shape == (0, 5)