diff --git a/src/ml_flashpoint/replication/replication_manager.py b/src/ml_flashpoint/replication/replication_manager.py index ca2b5df..48e6575 100644 --- a/src/ml_flashpoint/replication/replication_manager.py +++ b/src/ml_flashpoint/replication/replication_manager.py @@ -503,6 +503,7 @@ def shutdown(self): _LOGGER.info("Shutting down ReplicationManager and TransferService...") if self._transfer_service: self._transfer_service.shutdown() + self._transfer_service = None _LOGGER.info("TransferService shut down.") # TODO: Use the ip address return from c++ transfer service to avoid duplication. diff --git a/tests/replication/test_replication_manager.py b/tests/replication/test_replication_manager.py index 7820144..eb9e3dc 100644 --- a/tests/replication/test_replication_manager.py +++ b/tests/replication/test_replication_manager.py @@ -398,3 +398,17 @@ def test_async_replicate_single_node_skips(replication_manager, mocker): replication_manager._checkpoint_object_manager.close_buffer.assert_called_once_with( buffer_io, skip_close_if_symlink=True ) + + +def test_shutdown_clears_transfer_service(replication_manager): + """Tests that shutdown calls transfer_service.shutdown() and sets it to None.""" + # Given + mock_transfer_service = replication_manager._transfer_service + + # When + replication_manager.shutdown() + + # Then + mock_transfer_service.shutdown.assert_called_once() + + assert replication_manager._transfer_service is None