diff --git a/amqp/adapter/backpressure_handler.py b/amqp/adapter/backpressure_handler.py index e5d1b7b..97fd5a0 100644 --- a/amqp/adapter/backpressure_handler.py +++ b/amqp/adapter/backpressure_handler.py @@ -156,6 +156,7 @@ class BackpressureHandler: def start_backpressure_monitor(self) -> Thread: # Start the Backpressure monitor loop self.thread = Thread(target=self.backpressure_monitor_loop) + self.thread.daemon = True # This makes it a daemon thread self.thread.start() return self.thread @@ -185,10 +186,7 @@ class BackpressureHandler: IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 15, 90 if self.average_cpu_usage is None: self.average_cpu_usage = RunningAverage( - int( - max(_overload_duration_millis, _idle_duration_millis) - // _monitor_interval - ) + int(max(_overload_duration_millis, _idle_duration_millis) // _monitor_interval) ) _cpu_usage_state = CPU_ACTIVE _cpu_usage_changed = time.time() @@ -340,6 +338,4 @@ class BackpressureHandler: if BackpressureHandler._callback_list is not None: BackpressureHandler._callback_list["_wrap_rabbit_mq_api"] = 1 - await await_future( - asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api(), self.loop) - ) + await await_future(asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api(), self.loop)) diff --git a/tests/adapter/test_backpressure_handler.py b/tests/adapter/test_backpressure_handler.py index 51c981f..86d22d9 100644 --- a/tests/adapter/test_backpressure_handler.py +++ b/tests/adapter/test_backpressure_handler.py @@ -274,9 +274,7 @@ class TestBackpressureHandler: """ return AMQConfiguration("") - def test_backpressure_handler_initialization( - self, mock_channel, mock_loop, mock_config - ): + def test_backpressure_handler_initialization(self, mock_channel, mock_loop, mock_config): """ Test that the BackpressureHandler object is initialized correctly. """ @@ -334,14 +332,13 @@ class TestBackpressureHandler: Test the start_backpressure_monitor method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) + handler.do_loop = 1 # allow only 1 execution of the loop and then exit the thread, else tests will never end thread = handler.start_backpressure_monitor() assert isinstance(thread, Thread) assert handler.thread == thread @pytest.mark.asyncio - async def test_check_overload_condition_no_overload( - self, mock_channel, mock_loop, mock_config - ): + async def test_check_overload_condition_no_overload(self, mock_channel, mock_loop, mock_config): """ Test the check_overload_condition method when there is no overload. """ @@ -349,25 +346,19 @@ class TestBackpressureHandler: handler.current_parallel_executions = 1 handler.parallel_workers = 5 handler.last_backpressure_event = ScalingRequestAlert.UPDATE - with patch.object( - handler, "handle_backpressure_overload_event" - ) as mock_handle_overload: + with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload: await handler.check_overload_condition() mock_handle_overload.assert_not_called() @pytest.mark.asyncio - async def test_check_overload_condition_overload( - self, mock_channel, mock_loop, mock_config - ): + async def test_check_overload_condition_overload(self, mock_channel, mock_loop, mock_config): """ Test the check_overload_condition method when there is an overload. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.current_parallel_executions = 9 # Assuming parallel_workers is 10 handler.last_backpressure_event = ScalingRequestAlert.UPDATE - with patch.object( - handler, "handle_backpressure_overload_event" - ) as mock_handle_overload: + with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload: with patch("time.time") as mock_time: mock_time.return_value = 12345 await handler.check_overload_condition() @@ -385,9 +376,7 @@ class TestBackpressureHandler: handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.current_parallel_executions = 10 handler.last_backpressure_event = ScalingRequestAlert.OVERLOAD - with patch.object( - handler, "handle_backpressure_overload_event" - ) as mock_handle_overload: + with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload: await handler.check_overload_condition() mock_handle_overload.assert_not_called() @@ -396,9 +385,7 @@ class TestBackpressureHandler: # you might approach testing it, focusing on verifying that the correct methods # are called under specific conditions. This is NOT a complete, runnable test. @pytest.mark.asyncio - async def test_backpressure_monitor_loop( - self, mock_channel, mock_loop, mock_config - ): + async def test_backpressure_monitor_loop(self, mock_channel, mock_loop, mock_config): """ Test the backpressure_monitor_loop method (outline). """ @@ -406,9 +393,7 @@ class TestBackpressureHandler: handler.average_cpu_usage = RunningAverage( 100, initial_value=99 ) # Mock the RunningAverage in OVERLOAD state - handler.average_cpu_usage.is_filled = ( - True # Together with 0 values ensures OVERLOAD state - ) + handler.average_cpu_usage.is_filled = True # Together with 0 values ensures OVERLOAD state # Setup mocks for dependent methods handler.handle_backpressure_overload_event = AsyncMock() handler._handle_backpressure_idle_event = AsyncMock() @@ -428,9 +413,7 @@ class TestBackpressureHandler: handler.average_cpu_usage = RunningAverage( 100, initial_value=0 ) # Mock the RunningAverage in IDLE state - handler.average_cpu_usage.is_filled = ( - True # Together with 0 values ensures IDLE state - ) + handler.average_cpu_usage.is_filled = True # Together with 0 values ensures IDLE state handler.last_backpressure_event_time = 0 handler.last_data_message_time = 0 handler.do_loop = 1 @@ -441,9 +424,7 @@ class TestBackpressureHandler: handler._handle_backpressure_idle_event.assert_called_once() # 3. Simulate neither idle nor overload, but time_window passed - handler.average_cpu_usage = RunningAverage( - 100 - ) # Mock the RunningAverage in pritine state + handler.average_cpu_usage = RunningAverage(100) # Mock the RunningAverage in pritine state handler.last_backpressure_event_time = 0 handler.last_data_message_time = 0 handler.do_loop = 1 @@ -451,9 +432,7 @@ class TestBackpressureHandler: handler.publish_backpressure_request.assert_called_once() @pytest.mark.asyncio - async def test_handle_backpressure_overload_event( - self, mock_channel, mock_loop, mock_config - ): + async def test_handle_backpressure_overload_event(self, mock_channel, mock_loop, mock_config): """ Test the handle_backpressure_overload_event method. """ @@ -466,9 +445,7 @@ class TestBackpressureHandler: assert handler.last_data_message_time == 12345 @pytest.mark.asyncio - async def test_handle_backpressure_idle_event( - self, mock_channel, mock_loop, mock_config - ): + async def test_handle_backpressure_idle_event(self, mock_channel, mock_loop, mock_config): """ Test the _handle_backpressure_idle_event method. """ @@ -493,9 +470,7 @@ class TestBackpressureHandler: mock_async_run.return_value = asyncio.Future() mock_async_run.return_value.set_result(None) await handler.publish_backpressure_request( - ScaleRequestV1( - "test_service", "test_task", 100, 50, ScalingRequestAlert.UPDATE - ) + ScaleRequestV1("test_service", "test_task", 100, 50, ScalingRequestAlert.UPDATE) ) assert mock_async_run.call_count == 1 # handler.exchange.publish.assert_called_once() @@ -519,9 +494,7 @@ class TestBackpressureHandler: mock_exchange ) # Make it return the mock exchange await handler.publish_backpressure_request( - ScaleRequestV1( - "test_service", "test_task", 100, 50, ScalingRequestAlert.UPDATE - ) + ScaleRequestV1("test_service", "test_task", 100, 50, ScalingRequestAlert.UPDATE) ) assert mock_async_run.call_count == 2