From 6da72fed045575c227be5cd9e4b89f09c3375bbe Mon Sep 17 00:00:00 2001 From: Alex Golec Date: Tue, 25 Aug 2020 21:53:45 -0400 Subject: [PATCH] Fixes an issue with SSL contexts (#92) --- tda/streaming.py | 6 +++++- tests/streaming_test.py | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/tda/streaming.py b/tda/streaming.py index 41a5e8e..2075507 100644 --- a/tda/streaming.py +++ b/tda/streaming.py @@ -190,7 +190,11 @@ async def _init_from_principals(self, principals): # Initialize socket wss_url = 'wss://{}/ws'.format( principals['streamerInfo']['streamerSocketUrl']) - self._socket = await websockets.client.connect(wss_url, ssl=self._ssl_context) + if self._ssl_context: + self._socket = await websockets.client.connect( + wss_url, ssl=self._ssl_context) + else: + self._socket = await websockets.client.connect(wss_url) # Initialize miscellaneous parameters self._source = principals['streamerInfo']['appId'] diff --git a/tests/streaming_test.py b/tests/streaming_test.py index 495393e..a039d99 100644 --- a/tests/streaming_test.py +++ b/tests/streaming_test.py @@ -260,6 +260,42 @@ async def test_login_bad_response(self, ws_connect): with self.assertRaises(tda.streaming.UnexpectedResponseCode): await self.client.login() + @no_duplicates + @patch('tda.streaming.websockets.client.connect', new_callable=AsyncMock) + async def test_login_no_ssl_context(self, ws_connect): + self.client = StreamClient(self.http_client) + + self.http_client.get_user_principals.return_value = MockResponse( + account_principals(), True) + socket = AsyncMock() + ws_connect.return_value = socket + + socket.recv.side_effect = [json.dumps(self.success_response( + 0, 'ADMIN', 'LOGIN'))] + + await self.client.login() + + ws_connect.assert_awaited_once_with(ANY) + + + @no_duplicates + @patch('tda.streaming.websockets.client.connect', new_callable=AsyncMock) + async def test_login_ssl_context(self, ws_connect): + self.client = StreamClient(self.http_client, ssl_context='ssl_context') + + self.http_client.get_user_principals.return_value = MockResponse( + account_principals(), True) + socket = AsyncMock() + ws_connect.return_value = socket + + socket.recv.side_effect = [json.dumps(self.success_response( + 0, 'ADMIN', 'LOGIN'))] + + await self.client.login() + + ws_connect.assert_awaited_once_with(ANY, ssl='ssl_context') + + @no_duplicates @patch('tda.streaming.websockets.client.connect', new_callable=AsyncMock) async def test_login_unexpected_request_id(self, ws_connect):