Fix tests

dls-controls · Jun 6, 2016 · d948df2 · d948df2
1 parent cccf4a7
commit d948df2
Showing 1 changed file with 12 additions and 8 deletions.
diff --git a/tests/test_wscomms/test_wsservercomms.py b/tests/test_wscomms/test_wsservercomms.py
@@ -34,28 +34,32 @@ def test_setup(self, ioloop_mock, listen_mock):
 
         listen_mock.assert_called_once_with(1)
         self.assertEqual(ioloop_mock.current(), self.WS.loop)
-        self.assertEqual(self.WS.process,
-                         self.WS.WSApp.handlers[0][1][0].handler_class.process)
+        self.assertEqual(self.WS,
+                         self.WS.WSApp.handlers[0][1][0].handler_class.servercomms)
 
     @patch('malcolm.wscomms.wsservercomms.Application.listen')
     @patch('malcolm.wscomms.wsservercomms.IOLoop')
-    def test_start(self, ioloop_mock, _):
-        loop_mock = MagicMock()
-        ioloop_mock.current.return_value = loop_mock
+    def test_start(self, _, _2):
+        spawn_mock = MagicMock()
+        self.p.spawn.return_value = spawn_mock
+
         self.WS = WSServerComms("TestWebSocket", self.p, 1)
         self.WS.start_recv_loop()
 
-        loop_mock.start.assert_called_once_with()
+        self.assertEqual(spawn_mock, self.WS._loop_spawned)
 
     @patch('malcolm.wscomms.wsservercomms.Application.listen')
     @patch('malcolm.wscomms.wsservercomms.IOLoop')
     def test_stop(self, ioloop_mock, _):
         loop_mock = MagicMock()
         ioloop_mock.current.return_value = loop_mock
+
         self.WS = WSServerComms("TestWebSocket", self.p, 1)
+        self.WS._loop_spawned = MagicMock()
         self.WS.stop_recv_loop()
 
-        loop_mock.stop.assert_called_once_with()
+        loop_mock.add_callback.assert_called_once_with(loop_mock.stop)
+        self.WS._loop_spawned.wait.assert_called_once_with()
 
     @patch('malcolm.wscomms.wsservercomms.Request')
     @patch('malcolm.wscomms.wsservercomms.json')
@@ -77,7 +81,7 @@ def test_MWSH_on_message(self, _, _1, json_mock, request_mock):
             MWSH, "TestMessage")
 
         json_mock.loads.assert_called_once_with("TestMessage",
-                                                object_pairs_hook=OrderedDict())
+                                                object_pairs_hook=OrderedDict)
         request_mock.from_dict.assert_called_once_with(message_dict)
         self.p.handle_request.assert_called_once_with(request)