1 files changed, 15 insertions, 10 deletions

diff --git a/tests/test_swiftstory.py b/tests/test_ws_interface.py
index aac1865..6776334 100644
--- a/tests/test_swiftstory.py
+++ b/tests/test_ws_interface.py
@@ -1,56 +1,61 @@
 import unittest
 
-import swiftstory.__main__ as SwiftStory
+from swiftstory.interface.ws import WebsocketsInterface
 from swiftstory.status import error
 
 
-class TestSwiftStory(unittest.TestCase):
+class TestWebsocketsInterface(unittest.TestCase):
+    def setUp(self):
+        # XXX Passing None as the GameManager. Only OK if we don't try to join
+        # a game.
+        self.interface = WebsocketsInterface(None)
+
     def test_receive_invalid_json(self):
         self.assertEqual(
                 error("badly formatted json"),
-                SwiftStory.message_received_handler(client=None, message="{invalid_json}")
+                self.interface.message_received_handler(client=None, message="{invalid_json}")
                 )
 
     def test_receive_json_array(self):
         self.assertEqual(
                 error("invalid command"),
-                SwiftStory.message_received_handler(client=None, message='[]')
+                self.interface.message_received_handler(client=None, message='[]')
                 )
 
     def test_receive_json_number(self):
         self.assertEqual(
                 error("invalid command"),
-                SwiftStory.message_received_handler(client=None, message='2.3')
+                self.interface.message_received_handler(client=None, message='2.3')
                 )
 
     def test_receive_json_null(self):
         self.assertEqual(
                 error("invalid command"),
-                SwiftStory.message_received_handler(client=None, message='null')
+                self.interface.message_received_handler(client=None, message='null')
                 )
 
     def test_receive_unknown_command(self):
         self.assertEqual(
                 error("invalid command"),
-                SwiftStory.message_received_handler(client=None, message='{"op": "unknown"}')
+                self.interface.message_received_handler(client=None, message='{"op": "unknown"}')
                 )
 
     def test_receive_without_command(self):
         self.assertEqual(
                 error("invalid command"),
-                SwiftStory.message_received_handler(client=None, message='{}')
+                self.interface.message_received_handler(client=None, message='{}')
                 )
 
     def test_play_card_not_specified(self):
         payload = '{"op": "play_white_card"}'
         self.assertEqual(
                 error("field `card_id' is required"),
-                SwiftStory.message_received_handler(client=None, message=payload)
+                self.interface.message_received_handler(client=None, message=payload)
                 )
 
     def test_join_game_not_specified(self):
         payload = '{"op": "join_game"}'
         self.assertEqual(
                 error("field `game_name' is required"),
-                SwiftStory.message_received_handler(client=None, message=payload)
+                self.interface.message_received_handler(client=None, message=payload)
                 )