Fixes an issue with SSL contexts (#92)

alexgolec · Aug 26, 2020 · 6da72fe · 6da72fe
1 parent da090b2
commit 6da72fe
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 1 deletion.
diff --git 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
@@ -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):