add pytest timeout and annotate dense_image_warp test (#237)

Co-authored-by: Marton Havasi <[email protected]>
facebookresearch · Mar 18, 2024 · f207980 · f207980
1 parent 3437920
commit f207980
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 0 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -8,3 +8,6 @@ requires = [
 
 [tool.setuptools_scm]
 write_to = "neuralcompression/version.py"
+
+[tool.pytest.ini_options]
+timeout = 300
diff --git a/setup.cfg b/setup.cfg
@@ -80,6 +80,7 @@ dev =
     isort>=5.12.0
     mypy>=1.5.1
     pytest>=7.4.0
+    pytest-timeout>=2.3.1
 docs =
     myst-parser>=0.15.2
     sphinx-autodoc-typehints>=1.12.0
@@ -99,6 +100,7 @@ tests =
     opencv-python~=4.8.0.76
     pillow==9.4.0
     pytest==7.4.0
+    pytest-timeout==2.3.1
     pytorchvideo==0.1.5
     scipy==1.11.1
     torch==2.0.1

diff --git a/tests/functional/test_dense_image_warp.py b/tests/functional/test_dense_image_warp.py
@@ -5,11 +5,13 @@
 
 import pickle
 
+import pytest
 import torch
 
 from neuralcompression.functional import dense_image_warp
 
 
+@pytest.mark.timeout(method="thread")
 def test_dense_image_warp():
     with open("tests/cached_data/dense_image_warp.pkl", "rb") as f:
         data = pickle.load(f)