diff --git a/runtests.sh b/runtests.sh
index da9b3ed..8844f45 100755
--- a/runtests.sh
+++ b/runtests.sh
@@ -1,2 +1,2 @@
 #!/bin/sh
-python2.7 -m unittest discover py/tests -v
+python2.7 -m unittest discover tests -v
diff --git a/tests/test_Handshake.py b/tests/test_Handshake.py
index b0ac1c9..3ed0566 100644
--- a/tests/test_Handshake.py
+++ b/tests/test_Handshake.py
@@ -4,7 +4,7 @@
 import sys
 sys.path.insert(0, '..')
 
-from py.Handshake import Handshake
+from wifite.model.handshake import Handshake
 
 import unittest
 
diff --git a/tests/test_Target.py b/tests/test_Target.py
index d16f1d9..6d2f05f 100644
--- a/tests/test_Target.py
+++ b/tests/test_Target.py
@@ -1,7 +1,7 @@
 #!/usr/bin/python2.7
 # -*- coding: utf-8 -*-
 
-from py.Airodump import Airodump
+from wifite.tools.airodump import Airodump
 
 import unittest