diff --git a/tests/test_amq_service.py b/tests/test_amq_service.py new file mode 100644 index 0000000..ca97986 --- /dev/null +++ b/tests/test_amq_service.py @@ -0,0 +1,74 @@ +import asyncio +import unittest +from unittest.mock import AsyncMock, MagicMock, patch + +from amqp.service.amq_service import AMQService + + +class TestAMQService(unittest.IsolatedAsyncioTestCase): + def setUp(self): + # Mock the configuration + self.config = MagicMock() + self.config.amq_adapter.service_name = "test-service" + self.config.amq_adapter.swarm_task_id = "test-task" + self.config.amq_adapter.swarm_service_id = "test-service-id" + self.config.backpressure.threshold_load = 100 + + # Mock the service adapter + self.service_adapter = MagicMock() + + # Patch the get_unique_instance_id function + with patch("amqp.service.amq_service.get_unique_instance_id", return_value=12345): + # Create the AMQService instance + self.service = AMQService(self.config, self.service_adapter) + + # Mock the backpressure handler + self.service.backpressure_handler = MagicMock() + self.service.backpressure_handler.loop = asyncio.get_event_loop() + + async def test_backpressure_with_maximum(self): + """Test backpressure method with a specified maximum value.""" + # Call the backpressure method with a maximum value + self.service.backpressure(current_availability=50, maximum=200) + + # Check that update_backpressure_value was called with the correct values + self.service.backpressure_handler.update_backpressure_value.assert_called_once() + # Get the coroutine that was passed to run_coroutine_threadsafe + coro = self.service.backpressure_handler.update_backpressure_value.call_args[0] + self.assertEqual(coro, (50, 200)) + + async def test_backpressure_without_maximum(self): + """Test backpressure method without a maximum value.""" + # Call the backpressure method without a maximum value + self.service.backpressure(current_availability=50) + + # Check that update_backpressure_value was called with the correct values + self.service.backpressure_handler.update_backpressure_value.assert_called_once() + # Get the coroutine that was passed to run_coroutine_threadsafe + coro = self.service.backpressure_handler.update_backpressure_value.call_args[0] + self.assertEqual(coro, (50, 100)) # Should use threshold_load from config + + async def test_backpressure_with_negative_maximum(self): + """Test backpressure method with a negative maximum value.""" + # Call the backpressure method with a negative maximum value + self.service.backpressure(current_availability=50, maximum=-1) + + # Check that update_backpressure_value was called with the correct values + self.service.backpressure_handler.update_backpressure_value.assert_called_once() + # Get the coroutine that was passed to run_coroutine_threadsafe + coro = self.service.backpressure_handler.update_backpressure_value.call_args[0] + self.assertEqual(coro, (50, 100)) # Should use threshold_load from config + + async def test_backpressure_no_handler(self): + """Test backpressure method when no handler is available.""" + # Set the backpressure handler to None + self.service.backpressure_handler = None + + # Call the backpressure method + self.service.backpressure(current_availability=50) + + # Nothing should happen, no exception should be raised + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_backpressure_handler.py b/tests/test_backpressure_handler.py new file mode 100644 index 0000000..6fb5e14 --- /dev/null +++ b/tests/test_backpressure_handler.py @@ -0,0 +1,253 @@ +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()