Use the new InvalidStatus.code instead of checking str(exception).

python-websockets · Jul 22, 2017 · 1c8362f · 1c8362f
1 parent f76bb15
commit 1c8362f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/websockets/test_client_server.py b/websockets/test_client_server.py
@@ -272,9 +272,9 @@ def handshake(self, *args, **kwargs):
 
         self.start_server(klass=Protocol)
         try:
-            with self.assertRaises(InvalidHandshake) as cm:
+            with self.assertRaises(InvalidStatus) as cm:
                 self.start_client(path='invalid')
-            self.assertEqual(str(cm.exception), 'Bad status code: 404')
+            self.assertEqual(cm.exception.code, 404)
             self.assertFalse(State.handshake_called)
             # Check that our overridden handshake() is working correctly.
             self.start_client(path='valid')