import asyncio import time import unittest from unittest.mock import AsyncMock, MagicMock, patch from amqp.adapter.backpressure_handler import ( BackpressureHandler, RunningAverage, ScaleRequestV1, ScalingRequestAlert, ) class TestRunningAverage(unittest.TestCase): def test_running_average_empty(self): avg = RunningAverage(5) self.assertEqual(avg.get_average(), 0.0) def test_running_average_partial(self): avg = RunningAverage(5) avg.add(10) avg.add(20) self.assertEqual(avg.get_average(), 15.0) def test_running_average_full(self): avg = RunningAverage(3) avg.add(10) avg.add(20) avg.add(30) self.assertEqual(avg.get_average(), 20.0) def test_running_average_overflow(self): avg = RunningAverage(3) avg.add(10) avg.add(20) avg.add(30) avg.add(40) avg.add(50) self.assertEqual(avg.get_average(), 40.0) def test_get_current_values_partial(self): avg = RunningAverage(5) avg.add(10) avg.add(20) self.assertEqual(avg.get_current_values(), [10, 20]) def test_get_current_values_full(self): avg = RunningAverage(3) avg.add(10) avg.add(20) avg.add(30) self.assertEqual(avg.get_current_values(), [10, 20, 30]) def test_get_current_values_overflow(self): avg = RunningAverage(3) avg.add(10) avg.add(20) avg.add(30) avg.add(40) avg.add(50) self.assertEqual(avg.get_current_values(), [30, 40, 50]) class TestBackpressureHandler(unittest.IsolatedAsyncioTestCase): def setUp(self): self.channel = AsyncMock() self.loop = asyncio.get_event_loop() self.config = MagicMock() self.config.amq_adapter.swarm_service_id = "test-service" self.config.amq_adapter.swarm_task_id = "test-task" self.config.backpressure.threshold_load = 100 self.config.backpressure.time_window = 1 self.config.backpressure.idle_duration = 2 self.config.backpressure.cpu_monitoring_enabled = False self.config.backpressure.cpu_overload_duration = 1 # Record callbacks for testing BackpressureHandler._callback_list = {} self.handler = BackpressureHandler(self.channel, self.loop, self.config) self.handler.exchange = AsyncMock() # Set do_loop to 1 to run the loop only once self.handler.do_loop = 1 async def test_increase_decrease_current_load(self): self.handler.current_load = 0 self.handler.increase_current_load() self.assertEqual(self.handler.current_load, 1) self.handler.decrease_current_load() self.assertEqual(self.handler.current_load, 0) async def test_update_current_load(self): self.handler.current_load = 0 self.handler.update_current_load(5) self.assertEqual(self.handler.current_load, 5) async def test_update_last_data_message_time(self): old_time = self.handler.last_backpressure_event_time self.handler.update_last_data_message_time() self.assertGreater(self.handler.last_backpressure_event_time, old_time) @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request") async def test_update_backpressure_value(self, mock_publish): # Test updating with a higher maximum await self.handler.update_backpressure_value(50, 200) self.assertEqual(self.handler.current_load, 50) self.assertEqual(self.handler.maximum_load, 200) # Test updating with a lower maximum (should not change) await self.handler.update_backpressure_value(60, 50) self.assertEqual(self.handler.current_load, 60) self.assertEqual(self.handler.maximum_load, 200) @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request") async def test_check_overload_condition_normal(self, mock_publish): # Set up a normal state (60% available capacity) self.handler.maximum_load = 100 self.handler.current_load = 60 self.handler.last_backpressure_event = ScalingRequestAlert.UPDATE self.handler.last_backpressure_event_time = time.time() - 2 # 2 seconds ago await self.handler.check_overload_condition() # Should send an UPDATE event mock_publish.assert_called_once() args = mock_publish.call_args[0][0] self.assertEqual(args.requestType, ScalingRequestAlert.UPDATE) self.assertEqual(args.currentAvailability, 60) self.assertEqual(args.maxAvailability, 100) @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request") async def test_check_overload_condition_overload(self, mock_publish): # Set up an overload state (10% available capacity) self.handler.maximum_load = 100 self.handler.current_load = 10 self.handler.last_backpressure_event = ScalingRequestAlert.UPDATE self.handler.last_backpressure_event_time = time.time() - 2 # 2 seconds ago await self.handler.check_overload_condition() # Should send an OVERLOAD event mock_publish.assert_called_once() args = mock_publish.call_args[0][0] self.assertEqual(args.requestType, ScalingRequestAlert.OVERLOAD) self.assertEqual(args.currentAvailability, 10) self.assertEqual(args.maxAvailability, 100) @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request") async def test_check_overload_condition_idle(self, mock_publish): # Set up an idle state (90% available capacity) self.handler.maximum_load = 100 self.handler.current_load = 90 self.handler.last_backpressure_event = ScalingRequestAlert.IDLE self.handler.last_backpressure_event_time = time.time() - 3 # 3 seconds ago (> idle_duration) await self.handler.check_overload_condition() # Should send an IDLE event mock_publish.assert_called_once() args = mock_publish.call_args[0][0] self.assertEqual(args.requestType, ScalingRequestAlert.IDLE) self.assertEqual(args.currentAvailability, 90) self.assertEqual(args.maxAvailability, 100) @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request") async def test_check_overload_condition_auto_max(self, mock_publish): # Test with unset maximum that should be auto-detected self.handler.maximum_load = -1 self.handler.current_load = 50 self.handler.last_backpressure_event = ScalingRequestAlert.UPDATE self.handler.last_backpressure_event_time = time.time() - 2 # 2 seconds ago await self.handler.check_overload_condition() # Should set maximum to current and send UPDATE self.assertEqual(self.handler.maximum_load, 50) mock_publish.assert_called_once() args = mock_publish.call_args[0][0] self.assertEqual(args.requestType, ScalingRequestAlert.UPDATE) self.assertEqual(args.currentAvailability, 50) self.assertEqual(args.maxAvailability, 50) # Test with a higher value that should update the maximum mock_publish.reset_mock() await self.handler.update_backpressure_value(80, -1) # Should update maximum to 80 self.assertEqual(self.handler.maximum_load, 80) mock_publish.assert_called_once() args = mock_publish.call_args[0][0] self.assertEqual(args.currentAvailability, 80) self.assertEqual(args.maxAvailability, 80) @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request") async def test_handle_backpressure_overload_event(self, mock_publish): self.handler.maximum_load = 100 self.handler.current_load = 10 await self.handler.handle_backpressure_overload_event() mock_publish.assert_called_once() args = mock_publish.call_args[0][0] self.assertEqual(args.requestType, ScalingRequestAlert.OVERLOAD) self.assertEqual(args.currentAvailability, 10) self.assertEqual(args.maxAvailability, 100) @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request") async def test_handle_backpressure_idle_event(self, mock_publish): self.handler.maximum_load = 100 self.handler.current_load = 90 await self.handler._handle_backpressure_idle_event() mock_publish.assert_called_once() args = mock_publish.call_args[0][0] self.assertEqual(args.requestType, ScalingRequestAlert.IDLE) self.assertEqual(args.currentAvailability, 90) self.assertEqual(args.maxAvailability, 100) @patch("asyncio.run_coroutine_threadsafe") async def test_publish_backpressure_request(self, mock_run): scaling_request = ScaleRequestV1( "test-service", "test-task", 100, 50, ScalingRequestAlert.UPDATE ) await self.handler.publish_backpressure_request(scaling_request) # Check that run_coroutine_threadsafe was called self.assertEqual(mock_run.call_count, 1) # Check that the callback was recorded self.assertIn("_wrap_rabbit_mq_api", BackpressureHandler._callback_list) class TestScaleRequestV1(unittest.TestCase): def test_to_json(self): request = ScaleRequestV1( "test-service", "test-task", 100, 50, ScalingRequestAlert.UPDATE ) json_str = request.to_json() # Check that the JSON contains the expected fields self.assertIn('"version": 1', json_str) self.assertIn('"serviceId": "test-service"', json_str) self.assertIn('"taskId": "test-task"', json_str) self.assertIn('"maxAvailability": 100', json_str) self.assertIn('"currentAvailability": 50', json_str) self.assertIn(f'"requestType": "{ScalingRequestAlert.UPDATE.value}"', json_str) if __name__ == "__main__": unittest.main()