FIX: Live.block_for_close to call stop

Author: nmacholl

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## 0.28.1 - TBD
+
+#### Bug fixes
+- `Live.block_for_close` and `Live.wait_for_close` will now call `Live.stop` when a timeout is reached instead of `Live.terminate` to close the stream more gracefully
+
 ## 0.28.0 - 2024-02-01
 
 #### Enhancements

databento/live/client.py

@@ -538,6 +538,9 @@ def block_for_close(
         Block until the session closes or a timeout is reached. A session will
         close after `Live.stop` is called or the remote gateway disconnects.
 
+        If a `timeout` is specified, `Live.stop` will be called when the
+        timeout is reached.
+
         Parameters
         ----------
         timeout : float, optional
@@ -562,8 +565,8 @@ def block_for_close(
                 loop=Live._loop,
             ).result(timeout=timeout)
         except (futures.TimeoutError, KeyboardInterrupt) as exc:
-            logger.info("terminating session due to %s", type(exc).__name__)
-            self.terminate()
+            logger.info("closing session due to %s", type(exc).__name__)
+            self.stop()
             if isinstance(exc, KeyboardInterrupt):
                 raise
         except BentoError:
@@ -582,6 +585,9 @@ async def wait_for_close(
         session will close after `Live.stop` is called or the remote gateway
         disconnects.
 
+        If a `timeout` is specified, `Live.stop` will be called when the
+        timeout is reached.
+
         Parameters
         ----------
         timeout : float, optional
@@ -610,8 +616,8 @@ async def wait_for_close(
         try:
             await asyncio.wait_for(waiter, timeout=timeout)
         except (asyncio.TimeoutError, KeyboardInterrupt) as exc:
-            logger.info("terminating session due to %s", type(exc).__name__)
-            self.terminate()
+            logger.info("closing session due to %s", type(exc).__name__)
+            self.stop()
             if isinstance(exc, KeyboardInterrupt):
                 raise
         except BentoError:

tests/test_live_client.py

@@ -604,11 +604,10 @@ def test_live_block_for_close_timeout(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     """
-    Test that block_for_close terminates the session when the timeout is
-    reached.
+    Test that block_for_close stops the session when the timeout is reached.
     """
     # Arrange
-    monkeypatch.setattr(live_client, "terminate", MagicMock())
+    monkeypatch.setattr(live_client, "stop", MagicMock())
     live_client.subscribe(
         dataset=Dataset.GLBX_MDP3,
         schema=Schema.MBO,
@@ -619,7 +618,7 @@ def test_live_block_for_close_timeout(
 
     # Act, Assert
     live_client.block_for_close(timeout=0)
-    live_client.terminate.assert_called_once()  # type: ignore
+    live_client.stop.assert_called_once()  # type: ignore
 
 
 @pytest.mark.usefixtures("mock_live_server")
@@ -679,11 +678,10 @@ async def test_live_wait_for_close_timeout(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     """
-    Test that wait_for_close terminates the session when the timeout is
-    reached.
+    Test that wait_for_close stops the session when the timeout is reached.
     """
     # Arrange
-    monkeypatch.setattr(live_client, "terminate", MagicMock())
+    monkeypatch.setattr(live_client, "stop", MagicMock())
 
     # Act
     live_client.subscribe(
@@ -696,7 +694,7 @@ async def test_live_wait_for_close_timeout(
     await live_client.wait_for_close(timeout=0)
 
     # Assert
-    live_client.terminate.assert_called_once()  # type: ignore
+    live_client.stop.assert_called_once()  # type: ignore
 
 
 @pytest.mark.usefixtures("mock_live_server")