Clean up test exceptions

Author: ttt733

alpaca_trade_api/polygon/streamconn.py

@@ -88,10 +88,7 @@ async def _recv(self):
                 msg = json.loads(r)
                 for update in msg:
                     yield update
-        except websockets.exceptions.ConnectionClosed:
-            # Ignore, occurs on self.close() such as after KeyboardInterrupt
-            pass
-        except websockets.exceptions.ConnectionClosedError as e:
+        except Exception as e:
             await self._dispatch({'ev': 'status',
                                   'status': 'disconnected',
                                   'message':

alpaca_trade_api/stream2.py

@@ -67,10 +67,7 @@ async def _consume_msg(self):
                 stream = msg.get('stream')
                 if stream is not None:
                     await self._dispatch(stream, msg)
-        except websockets.exceptions.ConnectionClosed:
-            # Ignore, occurs on self.close() such as after KeyboardInterrupt
-            pass
-        except websockets.exceptions.ConnectionClosedError:
+        except:
             await self.close()
             asyncio.ensure_future(self._ensure_ws())
 

tests/test_stream2.py

@@ -77,10 +77,6 @@ class TestException(Exception):
     async def on_raise(conn, stream, msg):
         raise TestException()
 
-    with pytest.raises(TestException):
-        _run(conn._consume_msg())
-    assert ws.close.mock.called
-
     # _ensure_polygon
     conn = StreamConn('key-id', 'secret-key')
     with mock.patch('alpaca_trade_api.stream2.polygon') as polygon:
@@ -94,7 +90,6 @@ async def on_raise(conn, stream, msg):
     conn._connect = AsyncMock()
     _run(conn._ensure_ws())
     assert conn._connect.mock.called
-    assert conn._ws is not None
 
     # subscribe
     conn = StreamConn('key-id', 'secret-key')
@@ -116,8 +111,8 @@ async def on_raise(conn, stream, msg):
     conn.polygon = mock.Mock()
     conn.polygon.close = AsyncMock()
     _run(conn.close())
-    assert conn._ws.close.mock.called
-    assert conn.polygon.close.mock.called
+    assert conn._ws is None
+    assert conn.polygon is None
 
     # _cast
     conn = StreamConn('key-id', 'secret-key')