diff --git a/test/WsPrivateTest.py b/test/WsPrivateTest.py
index 2a8f0a0..af12427 100644
--- a/test/WsPrivateTest.py
+++ b/test/WsPrivateTest.py
@@ -1,7 +1,6 @@
 import time
 
-from WsPrivate import WsPrivate
-
+from okx.websocket.WsPrivate import WsPrivate
 
 def privateCallback(message):
     print("WsPrivate subscribe callback:", message)
diff --git a/test/WsPublicTest.py b/test/WsPublicTest.py
index 17c106e..573537d 100644
--- a/test/WsPublicTest.py
+++ b/test/WsPublicTest.py
@@ -1,7 +1,5 @@
 import time
-
-from WsPublic import WsPublic
-
+from okx.websocket.WsPublic import WsPublic
 
 def publicCallback(message):
     print("publicCallback", message)