diff --git a/tests/test_dragonnet.py b/tests/test_dragonnet.py
index 14eaabf6..1d4e613d 100644
--- a/tests/test_dragonnet.py
+++ b/tests/test_dragonnet.py
@@ -1,5 +1,5 @@
 from causalml.inference.tf import DragonNet
-from causalml.dataset.regression import *
+from causalml.dataset.regression import simulate_nuisance_and_easy_treatment
 import shutil