# # NOTE: need to install: pytest-asyncio # import asyncio import json from asyncio import AbstractEventLoop from threading import Thread from unittest.mock import AsyncMock, MagicMock, patch import pytest from aio_pika import Exchange from aio_pika.abc import AbstractRobustChannel from amqp.adapter.backpressure_handler import ( BackpressureHandler, RunningAverage, ScaleRequestV1, ) from amqp.config.amq_configuration import AMQConfiguration from amqp.model.model import ScalingRequestAlert class TestScaleRequestV1: """ A class containing pytest unit tests for the ScaleRequestV1 class. """ def test_scale_request_v1_initialization(self): """ Test that the ScaleRequestV1 object is initialized correctly. """ service_id = "test-service" task_id = "test-task" max_availability = 10 current_availability = 5 request_type = ScalingRequestAlert.UPDATE scale_request = ScaleRequestV1( service_id, task_id, max_availability, current_availability, request_type, ) assert scale_request.version == 1 assert scale_request.serviceId == service_id assert scale_request.taskId == task_id assert scale_request.max_availability == max_availability assert scale_request.current_availability == current_availability assert scale_request.requestType == request_type assert scale_request.thread is None def test_scale_request_v1_to_json(self): """ Test the to_json method of ScaleRequestV1. """ service_id = "test-service" task_id = "test-task" max_availability = 10 current_availability = 5 request_type = ScalingRequestAlert.UPDATE scale_request = ScaleRequestV1( service_id, task_id, max_availability, current_availability, request_type, ) json_output = scale_request.to_json() expected_json = json.dumps( { "version": 1, "serviceId": service_id, "taskId": task_id, "maxAvailability": max_availability, "currentAvailability": current_availability, "requestType": request_type.value, }, indent=2, ) assert json_output == expected_json def test_scale_request_v1_to_json_different_alert_type(self): """ Test the to_json method with a different ScalingRequestAlert type. """ service_id = "test-service" task_id = "test-task" max_availability = 10 current_availability = 5 request_type = ScalingRequestAlert.OVERLOAD # Change the alert type scale_request = ScaleRequestV1( service_id, task_id, max_availability, current_availability, request_type, ) json_output = scale_request.to_json() expected_json = json.dumps( { "version": 1, "serviceId": service_id, "taskId": task_id, "maxAvailability": max_availability, "currentAvailability": current_availability, "requestType": request_type.value, # Ensure the value is correctly serialized }, indent=2, ) assert json_output == expected_json class TestRunningAverage: """ A class containing pytest unit tests for the RunningAverage class. """ def test_running_average_initialization(self): """ Test that the RunningAverage object is initialized correctly. """ num_items = 5 ra = RunningAverage(num_items) assert ra.buffer == [0] * num_items assert ra.pointer == 0 assert ra.num_items == num_items assert ra.is_filled is False def test_running_average_add_one_value(self): """ Test adding a single value to the RunningAverage. """ ra = RunningAverage(3) ra.add(10) assert ra.buffer == [10, 0, 0] assert ra.pointer == 1 assert ra.is_filled is False def test_running_average_add_multiple_values_within_capacity(self): """ Test adding multiple values within the capacity of the RunningAverage buffer. """ ra = RunningAverage(3) ra.add(5) ra.add(10) ra.add(15) assert ra.buffer == [5, 10, 15] assert ra.pointer == 0 assert ra.is_filled is True def test_running_average_add_more_values_than_capacity(self): """ Test adding more values than the capacity of the RunningAverage buffer, checking wrap-around. """ ra = RunningAverage(3) ra.add(1) ra.add(2) ra.add(3) ra.add(4) # Overwrite the first value ra.add(5) # Overwrite the second value assert ra.buffer == [4, 5, 3] assert ra.pointer == 2 assert ra.is_filled is True def test_running_average_get_average_empty(self): """ Test get_average when no values have been added. """ ra = RunningAverage(3) assert ra.get_average() == 0.0 def test_running_average_get_average_one_value(self): """ Test get_average after adding one value. """ ra = RunningAverage(3) ra.add(10) assert ra.get_average() == 10.0 def test_running_average_get_average_multiple_values_within_capacity(self): """ Test get_average with multiple values added within the buffer's capacity. """ ra = RunningAverage(3) ra.add(1) ra.add(2) ra.add(3) assert ra.get_average() == 2.0 def test_running_average_get_average_more_values_than_capacity(self): """ Test get_average after adding more values than the buffer's capacity (check wrap-around). """ ra = RunningAverage(3) ra.add(1) ra.add(2) ra.add(3) ra.add(4) ra.add(5) assert ra.get_average() == 4.0 def test_get_current_values_empty(self): """Test get_current_values when no values have been added.""" ra = RunningAverage(3) assert ra.get_current_values() == [] def test_get_current_values_one_value(self): """Test get_current_values after adding one value.""" ra = RunningAverage(3) ra.add(10) assert ra.get_current_values() == [10] def test_get_current_values_multiple_values_within_capacity(self): """Test get_current_values with multiple values within capacity.""" ra = RunningAverage(3) ra.add(1) ra.add(2) ra.add(3) assert ra.get_current_values() == [1, 2, 3] def test_get_current_values_more_values_than_capacity(self): """Test get_current_values after adding more values than capacity (wrap-around).""" ra = RunningAverage(3) ra.add(1) ra.add(2) ra.add(3) ra.add(4) ra.add(5) assert ra.get_current_values() == [3, 4, 5] def test_get_current_values_wrap_around(self): """Test get_current_values when the buffer has wrapped around.""" ra = RunningAverage(4) ra.add(1) ra.add(2) ra.add(3) ra.add(4) ra.add(5) # Overwrite 1 ra.add(6) # Overwrite 2 assert ra.get_current_values() == [3, 4, 5, 6] # Test for BackpressureHandler class TestBackpressureHandler: """ A class containing pytest unit tests for the BackpressureHandler class. """ @pytest.fixture def mock_channel(self): """ Pytest fixture to create a mock AbstractRobustChannel. """ return AsyncMock(spec=AbstractRobustChannel) @pytest.fixture def mock_loop(self): """ Pytest fixture to create a mock AbstractEventLoop. """ return MagicMock(spec=AbstractEventLoop) @pytest.fixture def mock_config(self): """ Pytest fixture to create a mock AMQConfiguration. """ amq_cfg = AMQConfiguration("") amq_cfg.backpressure.cpu_monitoring_enabled = True return amq_cfg def test_backpressure_handler_initialization(self, mock_channel, mock_loop, mock_config): """ Test that the BackpressureHandler object is initialized correctly. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) assert handler.channel == mock_channel assert handler.loop == mock_loop assert handler.config == mock_config assert handler.swarm_service_id == "clever-amqp-service" assert handler.swarm_task_id == "python-adapter" assert handler.max_availability == mock_config.backpressure.threshold_load assert handler.average_cpu_usage is None def test_increase_current_load(self, mock_channel, mock_loop, mock_config): """ Test the increase_current_load method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) initial_load = handler.current_availability handler.increase_current_load() assert handler.current_availability == initial_load + 1 def test_decrease_current_load(self, mock_channel, mock_loop, mock_config): """ Test the decrease_current_load method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.current_availability = 5 handler.decrease_current_load() assert handler.current_availability == 4 handler.current_availability = 0 handler.decrease_current_load() assert handler.current_availability == 0 # Should not go below 0 def test_update_current_load(self, mock_channel, mock_loop, mock_config): """ Test the update_current_availability method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.update_current_availability(100) assert handler.current_availability == 100 def test_update_last_data_message_time(self, mock_channel, mock_loop, mock_config): """ Test the update_last_backpressure_event_time method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) with patch("time.time") as mock_time: mock_time.return_value = 12345 handler.update_last_backpressure_event_time() assert handler.last_backpressure_event_time == 12345 def test_start_backpressure_monitor(self, mock_channel, mock_loop, mock_config): """ 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): """ Test the check_overload_condition method when there is no overload and only short time elapsed from last message """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.current_availability = 1 handler.max_availability = 5 handler.last_backpressure_event = ScalingRequestAlert.UPDATE handler.publish_backpressure_request = AsyncMock() handler.last_backpressure_event_time = 12340 with patch("time.time") as mock_time: mock_time.return_value = 12345 with patch.object( handler, "handle_backpressure_overload_event" ) as mock_handle_overload: await handler.check_overload_condition() mock_handle_overload.assert_not_called() if ( 12345 - handler.last_backpressure_event_time > mock_config.backpressure.time_window ): handler.publish_backpressure_request.assert_called_once() assert handler.last_backpressure_event == ScalingRequestAlert.UPDATE assert handler.last_backpressure_event_time == 12345 else: handler.publish_backpressure_request.assert_not_called() @pytest.mark.asyncio 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_availability = 9 # Assuming maxAvailability is 10 handler.last_backpressure_event = ScalingRequestAlert.UPDATE handler.publish_backpressure_request = AsyncMock() with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload: mock_handle_overload.return_value = None with patch("time.time") as mock_time: mock_time.return_value = 12345 await handler.check_overload_condition() mock_handle_overload.assert_called_once() assert handler.last_backpressure_event_time == 12345 assert handler.last_backpressure_event == ScalingRequestAlert.OVERLOAD @pytest.mark.asyncio async def test_check_overload_condition_already_overload( self, mock_channel, mock_loop, mock_config ): """ Test the check_overload_condition method when the last event was already an overload. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.current_availability = 10 handler.last_backpressure_event = ScalingRequestAlert.OVERLOAD handler.publish_backpressure_request = AsyncMock() with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload: await handler.check_overload_condition() mock_handle_overload.assert_not_called() # backpressure_monitor_loop is difficult to test directly without significant mocking # and potentially long-running tests. The following is a basic outline of how # 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): """ Test the backpressure_monitor_loop method (outline). """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) 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 # Setup mocks for dependent methods handler.handle_backpressure_overload_event = AsyncMock() handler._handle_backpressure_idle_event = AsyncMock() handler.publish_backpressure_request = AsyncMock() # 1. Simulate an overload condition: handler.last_backpressure_event_time = 0 # Force event trigger handler.last_backpressure_event_time = 0 handler.do_loop = 1 handler.config.backpressure.cpu_overload_duration = ( -1 ) # Force overload condition even for recent change await handler._backpressure_monitor() handler.handle_backpressure_overload_event.assert_called_once() # 2. Simulate an idle condition 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.last_backpressure_event_time = 0 handler.last_backpressure_event_time = 0 handler.do_loop = 1 handler.config.backpressure.idle_duration = ( -1 ) # Force idle condition even for recent change await handler._backpressure_monitor() handler._handle_backpressure_idle_event.assert_called_once() # 3. Simulate neither idle nor overload, but time_window passed handler.handle_backpressure_overload_event.call_count = 0 handler._handle_backpressure_idle_event.call_count = 0 handler.average_cpu_usage = RunningAverage(100) # Mock the RunningAverage in pritine state handler.last_backpressure_event_time = 0 handler.last_backpressure_event_time = 0 handler.do_loop = 1 await handler._backpressure_monitor() assert ( handler.publish_backpressure_request.call_count + handler.handle_backpressure_overload_event.call_count + handler._handle_backpressure_idle_event.call_count ) > 0 @pytest.mark.asyncio async def test_handle_backpressure_overload_event(self, mock_channel, mock_loop, mock_config): """ Test the handle_backpressure_overload_event method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.publish_backpressure_request = AsyncMock() with patch("time.time") as mock_time: mock_time.return_value = 12345 await handler.handle_backpressure_overload_event() assert handler.publish_backpressure_request.call_count == 1 @pytest.mark.asyncio async def test_handle_backpressure_idle_event(self, mock_channel, mock_loop, mock_config): """ Test the _handle_backpressure_idle_event method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.publish_backpressure_request = AsyncMock() with patch("time.time") as mock_time: mock_time.return_value = 12345 await handler._handle_backpressure_idle_event() assert handler.publish_backpressure_request.call_count == 1 @pytest.mark.asyncio async def test_update_backpressure_value_overload(self, mock_channel, mock_loop, mock_config): """ Test the update_backpressure_value method when resource usage is in OVERLOAD state. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.handle_backpressure_overload_event = AsyncMock() handler._handle_backpressure_idle_event = AsyncMock() handler.publish_backpressure_request = AsyncMock() # Set up initial state handler.last_backpressure_event = ScalingRequestAlert.UPDATE handler.last_backpressure_event_time = 10 # time of last message handler._resource_usage_changed = 10 # time of last change with patch("time.time") as mock_time: mock_time.return_value = 100 # Call with resource usage above OVERLOAD threshold (90% of maximum) await handler.update_backpressure_value(95, 100) # Debug output print(f"Last backpressure event: {handler.last_backpressure_event}") print(f"Last backpressure event time: {handler.last_backpressure_event_time}") print(f"Time window: {mock_config.backpressure.time_window}") print("Resource usage: 95, Maximum: 100") print(f"Overload threshold calculation: {round(0.9 * 100)}") # Verify OVERLOAD event was triggered handler.handle_backpressure_overload_event.assert_called_once() assert handler.last_backpressure_event == ScalingRequestAlert.OVERLOAD assert handler.last_backpressure_event_time == 100 @pytest.mark.asyncio async def test_update_backpressure_value_idle(self, mock_channel, mock_loop, mock_config): """ Test the update_backpressure_value method when resource usage is in IDLE state. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler._handle_backpressure_idle_event = AsyncMock() handler.handle_backpressure_overload_event = AsyncMock() handler.publish_backpressure_request = AsyncMock() # Set up initial state handler.last_backpressure_event = ScalingRequestAlert.UPDATE handler.last_backpressure_event_time = 0 handler._resource_usage_changed = 0 with patch("time.time") as mock_time: # First call to set state to IDLE mock_time.return_value = 50 await handler.update_backpressure_value(0, 10) # Second call after idle_duration has passed mock_time.return_value = 50 + mock_config.backpressure.idle_duration + 1 await handler.update_backpressure_value(0, 10) # Verify IDLE event was triggered handler._handle_backpressure_idle_event.assert_called_once() assert handler.last_backpressure_event == ScalingRequestAlert.IDLE assert ( handler.last_backpressure_event_time == 50 + mock_config.backpressure.idle_duration + 1 ) @pytest.mark.asyncio async def test_update_backpressure_value_update(self, mock_channel, mock_loop, mock_config): """ Test the update_backpressure_value method when resource usage is in ACTIVE state. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.publish_backpressure_request = AsyncMock() handler.handle_backpressure_overload_event = AsyncMock() handler._handle_backpressure_idle_event = AsyncMock() # Set up initial state handler.last_backpressure_event = ScalingRequestAlert.IDLE handler.last_backpressure_event_time = 0 with patch("time.time") as mock_time: mock_time.return_value = 100 # Call with resource usage between thresholds await handler.update_backpressure_value(50, 100) # Verify UPDATE event was triggered assert handler.publish_backpressure_request.call_count == 1 assert handler.last_backpressure_event == ScalingRequestAlert.UPDATE assert handler.last_backpressure_event_time == 100 @pytest.mark.asyncio async def test_publish_backpressure_request_exchange_exists( self, mock_channel, mock_loop, mock_config ): """ Test the publish_backpressure_request method when the exchange already exists. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.exchange = AsyncMock(spec=Exchange) with patch("asyncio.run_coroutine_threadsafe") as mock_async_run: 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) ) assert mock_async_run.call_count == 1 # handler.exchange.publish.assert_called_once() @pytest.mark.asyncio async def test_publish_backpressure_request_exchange_does_not_exist( self, mock_channel, mock_loop, mock_config ): """ Test the publish_backpressure_request method when the exchange does not exist and needs to be created. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.exchange = None # Simulate exchange not existing mock_exchange = AsyncMock(spec=Exchange) mock_channel.get_exchange.return_value = mock_exchange BackpressureHandler._callback_list = {} with patch("asyncio.run_coroutine_threadsafe") as mock_async_run: mock_async_run.return_value = asyncio.Future() mock_async_run.return_value.set_result( mock_exchange ) # Make it return the mock exchange await handler.publish_backpressure_request( ScaleRequestV1("test_service", "test_task", 100, 50, ScalingRequestAlert.UPDATE) ) assert mock_async_run.call_count == 2 assert "_wrap_rabbit_mq_api_init" in BackpressureHandler._callback_list assert "_wrap_rabbit_mq_api" in BackpressureHandler._callback_list @pytest.mark.asyncio async def test_update_backpressure_value(self, mock_channel, mock_loop, mock_config): """ Test the update_backpressure_value method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.update_current_availability = MagicMock() handler.update_last_backpressure_event_time = MagicMock() handler.check_overload_condition = AsyncMock() # Test with current_availability value and same maximum await handler.update_backpressure_value(5, handler.max_availability) handler.update_current_availability.assert_called_once_with(5) handler.update_last_backpressure_event_time.assert_called_once() handler.check_overload_condition.assert_called_once() # Reset mocks handler.update_current_availability.reset_mock() handler.update_last_backpressure_event_time.reset_mock() handler.check_overload_condition.reset_mock() # Test with current_availability value and different maximum original_workers = handler.max_availability await handler.update_backpressure_value(3, original_workers + 5) handler.update_current_availability.assert_called_once_with(3) handler.update_last_backpressure_event_time.assert_called_once() handler.check_overload_condition.assert_called_once() assert handler.max_availability == original_workers + 5 @pytest.mark.asyncio async def test_check_overload_condition_idle(self, mock_channel, mock_loop, mock_config): """ Test the check_overload_condition method when service is idle. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler._handle_backpressure_idle_event = AsyncMock() # Set up initial state for IDLE condition handler.current_availability = 0 handler.last_backpressure_event = ScalingRequestAlert.IDLE with patch("time.time") as mock_time: # Set current_availability time current_time = 1000 mock_time.return_value = current_time # Set last data message time to be older than idle_duration handler.last_backpressure_event_time = ( current_time - mock_config.backpressure.idle_duration - 10 ) # Call the method await handler.check_overload_condition() # Verify IDLE event was triggered handler._handle_backpressure_idle_event.assert_called_once() assert handler.last_backpressure_event == ScalingRequestAlert.IDLE assert handler.last_backpressure_event_time == current_time @pytest.mark.asyncio async def test_check_overload_condition_update(self, mock_channel, mock_loop, mock_config): """ Test the check_overload_condition method when service is in normal state (UPDATE). """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.publish_backpressure_request = AsyncMock() # Set up initial state for UPDATE condition handler.current_availability = 2 handler.max_availability = 10 handler.last_backpressure_event = ScalingRequestAlert.IDLE # Coming from IDLE state with patch("time.time") as mock_time: # Set current_availability time current_time = 1000 mock_time.return_value = current_time # Set last event time to be older than time_window handler.last_backpressure_event_time = ( current_time - mock_config.backpressure.time_window - 10 ) # Call the method await handler.check_overload_condition() # Verify UPDATE event was triggered handler.publish_backpressure_request.assert_called_once() assert handler.last_backpressure_event == ScalingRequestAlert.UPDATE assert handler.last_backpressure_event_time == current_time def test_start_backpressure_monitor_disabled(self, mock_channel, mock_loop, mock_config): """ Test the start_backpressure_monitor method when CPU monitoring is disabled. """ # Set CPU monitoring to disabled in config mock_config.backpressure.cpu_monitoring_enabled = False handler = BackpressureHandler(mock_channel, mock_loop, mock_config) thread = handler.start_backpressure_monitor() # Verify no thread was started assert thread is not None assert hasattr(handler, "thread") or handler.thread is not None def test_start_backpressure_monitor_enabled(self, mock_channel, mock_loop, mock_config): """ Test the start_backpressure_monitor method when CPU monitoring is enabled. """ # Set CPU monitoring to enabled in config mock_config.backpressure.cpu_monitoring_enabled = True handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.do_loop = 1 # allow only 1 execution of the loop and then exit the thread thread = handler.start_backpressure_monitor() # Verify thread was started assert isinstance(thread, Thread) assert handler.thread == thread # Since I need also to mockup the asyncio.run_coroutine_threadsafe, the following mockups are # not called because mockup overrides that call with preset return value. # mock_channel.get_exchange.assert_called_once_with(name="cleverthis.clevermicro.management") # handler.exchange.publish.assert_called_once()