Stanislav Hejny
336 lines
14 KiB
Python
336 lines
14 KiB
Python
import asyncio
|
|
import time
|
|
import unittest
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from amqp.adapter.backpressure_handler import (
|
|
BackpressureHandler,
|
|
ScaleRequestV1,
|
|
ScalingRequestAlert,
|
|
)
|
|
|
|
|
|
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.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_availability = 0
|
|
self.handler.update_current_availability(5)
|
|
self.assertEqual(self.handler.current_availability, 5)
|
|
|
|
async def test_update_last_data_message_time(self):
|
|
old_time = self.handler.last_backpressure_event_time
|
|
self.handler.update_last_backpressure_event_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_availability, 50)
|
|
self.assertEqual(self.handler.max_availability, 200)
|
|
|
|
# Test updating with a lower maximum (should not change)
|
|
await self.handler.update_backpressure_value(60, 50)
|
|
self.assertEqual(self.handler.current_availability, 60)
|
|
self.assertEqual(self.handler.max_availability, 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.max_availability = 100
|
|
self.handler.current_availability = 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.current_availability, 60)
|
|
self.assertEqual(args.max_availability, 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.max_availability = 100
|
|
self.handler.current_availability = 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.current_availability, 10)
|
|
self.assertEqual(args.max_availability, 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.max_availability = 100
|
|
self.handler.current_availability = 90
|
|
self.handler.last_backpressure_event = ScalingRequestAlert.IDLE
|
|
# emulate fact that IDLE state is for 3 seconds longer than required minimum for sending IDLE event
|
|
self.handler.last_backpressure_event_time = (
|
|
time.time() - self.config.backpressure.idle_duration - 3
|
|
)
|
|
|
|
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.current_availability, 90)
|
|
self.assertEqual(args.max_availability, 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.max_availability = -1
|
|
self.handler.current_availability = 50
|
|
self.handler.last_backpressure_event = ScalingRequestAlert.UPDATE
|
|
self.handler.last_backpressure_event_time = (
|
|
time.time() - self.config.backpressure.time_window - 2
|
|
)
|
|
|
|
await self.handler.check_overload_condition()
|
|
|
|
# Should set maximum to current_availability and send UPDATE
|
|
# self.assertEqual(self.handler.max_availability, 50)
|
|
mock_publish.assert_called_once()
|
|
args = mock_publish.call_args[0][0]
|
|
self.assertEqual(args.requestType, ScalingRequestAlert.UPDATE)
|
|
self.assertEqual(args.current_availability, 50)
|
|
self.assertEqual(args.max_availability, -1)
|
|
|
|
# Test with a higher value that should update the maximum
|
|
mock_publish.reset_mock()
|
|
await self.handler.update_backpressure_value(80, 80)
|
|
|
|
# Should update maximum to 80
|
|
self.assertEqual(self.handler.max_availability, 80)
|
|
# mock_publish.assert_called_once()
|
|
# args = mock_publish.call_args[0][0]
|
|
# self.assertEqual(args.current_availability, 80)
|
|
# self.assertEqual(args.max_availability, 80)
|
|
|
|
@patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
|
|
async def test_handle_backpressure_overload_event(self, mock_publish):
|
|
self.handler.max_availability = 100
|
|
self.handler.current_availability = 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.current_availability, 10)
|
|
self.assertEqual(args.max_availability, 100)
|
|
|
|
@patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
|
|
async def test_handle_backpressure_idle_event(self, mock_publish):
|
|
self.handler.max_availability = 100
|
|
self.handler.current_availability = 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.current_availability, 90)
|
|
self.assertEqual(args.max_availability, 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)
|
|
|
|
@patch("amqp.adapter.backpressure_handler.BackpressureHandler.check_overload_condition")
|
|
@patch("asyncio.sleep")
|
|
async def test_backpressure_monitor(self, mock_sleep, mock_check_overload):
|
|
"""Test the _backpressure_monitor method"""
|
|
# Set do_loop to 3 to run the loop 3 times
|
|
self.handler.do_loop = 3
|
|
|
|
await self.handler._backpressure_monitor()
|
|
|
|
# Should call check_overload_condition 3 times
|
|
self.assertEqual(mock_check_overload.call_count, 3)
|
|
|
|
# Should call sleep 3 times with 0.5 second interval
|
|
self.assertEqual(mock_sleep.call_count, 3)
|
|
for call in mock_sleep.call_args_list:
|
|
self.assertEqual(call[0][0], 0.5)
|
|
|
|
@patch("asyncio.new_event_loop")
|
|
def test_backpressure_monitor_loop(self, mock_new_loop):
|
|
"""Test the backpressure_monitor_loop method"""
|
|
mock_loop = MagicMock()
|
|
mock_new_loop.return_value = mock_loop
|
|
|
|
# Set do_loop to 2 to run the loop 2 times
|
|
self.handler.do_loop = 2
|
|
|
|
self.handler.backpressure_monitor_loop()
|
|
|
|
# Should create a new event loop
|
|
mock_new_loop.assert_called_once()
|
|
|
|
# Should run the _backpressure_monitor coroutine on the loop instance
|
|
mock_loop.run_until_complete.assert_called_once()
|
|
|
|
# Check that the coroutine passed to run_until_complete is _backpressure_monitor
|
|
call_args = mock_loop.run_until_complete.call_args[0][0]
|
|
# The coroutine should be the result of calling _backpressure_monitor
|
|
self.assertTrue(asyncio.iscoroutine(call_args))
|
|
|
|
@patch("amqp.adapter.backpressure_handler.BackpressureHandler.check_overload_condition")
|
|
@patch("asyncio.sleep")
|
|
async def test_backpressure_monitor_infinite_loop(self, mock_sleep, mock_check_overload):
|
|
"""Test the _backpressure_monitor method with infinite loop (do_loop = -1)"""
|
|
# Set do_loop to -1 for infinite loop
|
|
self.handler.do_loop = -1
|
|
|
|
# Mock sleep to raise an exception after a few calls to break the infinite loop
|
|
call_count = 0
|
|
|
|
def mock_sleep_side_effect(*args):
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count > 3: # Break after 3 iterations
|
|
raise Exception("Test break")
|
|
return None
|
|
|
|
mock_sleep.side_effect = mock_sleep_side_effect
|
|
|
|
# Should raise the exception we injected
|
|
with self.assertRaises(Exception):
|
|
await self.handler._backpressure_monitor()
|
|
|
|
# Should call check_overload_condition 4 times before breaking
|
|
self.assertEqual(mock_check_overload.call_count, 4)
|
|
|
|
# Should call sleep 4 times
|
|
self.assertEqual(mock_sleep.call_count, 4)
|
|
|
|
@patch("amqp.adapter.backpressure_handler.BackpressureHandler.check_overload_condition")
|
|
@patch("asyncio.sleep")
|
|
async def test_backpressure_monitor_zero_loop(self, mock_sleep, mock_check_overload):
|
|
"""Test the _backpressure_monitor method with do_loop = 0 (should not run)"""
|
|
# Set do_loop to 0 to not run the loop
|
|
self.handler.do_loop = 0
|
|
|
|
await self.handler._backpressure_monitor()
|
|
|
|
# Should not call check_overload_condition
|
|
mock_check_overload.assert_not_called()
|
|
|
|
# Should not call sleep
|
|
mock_sleep.assert_not_called()
|
|
|
|
def test_do_loop_helper(self):
|
|
"""Test the _do_loop helper method"""
|
|
# Test with positive value
|
|
self.handler.do_loop = 3
|
|
self.assertTrue(self.handler._do_loop()) # Should return True and decrement
|
|
self.assertEqual(self.handler.do_loop, 2)
|
|
|
|
self.assertTrue(self.handler._do_loop()) # Should return True and decrement
|
|
self.assertEqual(self.handler.do_loop, 1)
|
|
|
|
self.assertTrue(self.handler._do_loop()) # Should return True and decrement
|
|
self.assertEqual(self.handler.do_loop, 0)
|
|
|
|
self.assertFalse(self.handler._do_loop()) # Should return False when do_loop is 0
|
|
|
|
# Test with negative value (infinite loop)
|
|
self.handler.do_loop = -1
|
|
self.assertTrue(self.handler._do_loop()) # Should return True and not change
|
|
self.assertEqual(self.handler.do_loop, -1)
|
|
|
|
# Test with zero
|
|
self.handler.do_loop = 0
|
|
self.assertFalse(self.handler._do_loop()) # Should return False
|
|
|
|
@patch("amqp.adapter.backpressure_handler.Thread")
|
|
def test_start_backpressure_monitor(self, mock_thread):
|
|
"""Test the start_backpressure_monitor method"""
|
|
mock_thread_instance = MagicMock()
|
|
mock_thread.return_value = mock_thread_instance
|
|
|
|
result = self.handler.start_backpressure_monitor()
|
|
|
|
# Should create a Thread with the correct target
|
|
mock_thread.assert_called_once_with(target=self.handler.backpressure_monitor_loop)
|
|
|
|
# Should set daemon to True
|
|
mock_thread_instance.daemon = True
|
|
|
|
# Should start the thread
|
|
mock_thread_instance.start.assert_called_once()
|
|
|
|
# Should return the thread instance
|
|
self.assertEqual(result, mock_thread_instance)
|
|
|
|
# Should store the thread reference
|
|
self.assertEqual(self.handler.thread, mock_thread_instance)
|
|
|
|
|
|
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()
|