diff --git a/test/unit/test_segmentation.py b/test/unit/test_segmentation.py
index cc97b88..c195e0c 100644
--- a/test/unit/test_segmentation.py
+++ b/test/unit/test_segmentation.py
@@ -1,5 +1,6 @@
 import numpy as np
 from numpy.testing import assert_allclose
+import pytest
 import matplotlib.pyplot as plt
 from dpm_tools.segmentation import seeded_region_growing as srg, statistical_region_merging as srm
 # import pyvista as pv