diff --git a/tests/test_API.py b/tests/test_API.py
index 2b1ba4c..77b8700 100644
--- a/tests/test_API.py
+++ b/tests/test_API.py
@@ -10,8 +10,8 @@
 import lshmm as ls
 import lshmm.forward_backward.fb_diploid as fbd
 import lshmm.forward_backward.fb_haploid as fbh
-import lshmm.vit_diploid_variants_samples as vd
-import lshmm.vit_haploid_variants_samples as vh
+import lshmm.vit_diploid as vd
+import lshmm.vit_haploid as vh
 
 EQUAL_BOTH_HOM = 4
 UNEQUAL_BOTH_HOM = 0