test: add unit tests for BackpressureHandler and AMQService classes
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
This commit is contained in:
Stanislav Hejny
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user