Add duration parameter to DataStream run method

alpaca/data/live/websocket.py

@@ -3,6 +3,7 @@
 import queue
 from collections import defaultdict
 from typing import Callable, Dict, List, Optional, Tuple, Union
+from datetime import timedelta
 
 import msgpack
 import websockets
@@ -359,10 +360,24 @@ async def _run_forever(self) -> None:
             finally:
                 await asyncio.sleep(0)
 
-    def run(self) -> None:
-        """Starts up the websocket connection's event loop"""
+    async def _run(self, duration: timedelta = None) -> None:
+        timeout_seconds = duration.total_seconds() if duration != None else None
+        run_task = asyncio.create_task(self._run_forever())
+        await asyncio.wait([run_task], timeout=timeout_seconds)
+
+    def run(self, duration: timedelta = None) -> None:
+        """Starts up the websocket connection's event loop
+
+        Parameters:
+        -----------
+        duration: timedelta, default 'None'
+            Duration of event loop before timeout."""
         try:
-            asyncio.run(self._run_forever())
+            asyncio.run(self._run(duration))
+        except TypeError as e:
+            print(e)
+            print("invalid duration type entered")
+            pass
         except KeyboardInterrupt:
             print("keyboard interrupt, bye")
             pass

tests/data/test_websockets.py

@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timedelta
 
 import pytest
 from msgpack.ext import Timestamp
@@ -232,3 +232,18 @@ async def handler_star(d):
     assert len(articles_b) == 1
     assert len(articles_star) == 2
     assert articles_star[1].headline == "c"
+
+
+@pytest.mark.parametrize("duration_seconds", [(1), (2), (3)])
+def test_run(ws_client: DataStream, duration_seconds: int) -> None:
+    """Testing for different durations.
+
+    Parameters:
+    -----------
+    ws_client (DataStream): DataStream to test.
+    duration_seconds (int): Duration to test.
+    """
+    start_time = datetime.now()
+    duration = timedelta(seconds=duration_seconds)
+    ws_client.run(duration)
+    assert round((datetime.now() - start_time).total_seconds()) >= duration_seconds