diff --git a/test/integration/test_drainage.py b/test/integration/test_drainage.py
index 9b7bfd6e0..71179f43c 100644
--- a/test/integration/test_drainage.py
+++ b/test/integration/test_drainage.py
@@ -1,7 +1,10 @@
 import numpy as np
 import porespy as ps
 import matplotlib.pyplot as plt
-from pyedt import edt
+try:
+    from pyedt import edt
+except ModuleNotFoundError:
+    from edt import edt
 
 
 def test_drainage():