# # 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.maxAvailability == max_availability assert scale_request.currentAvailability == 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_string. """ service_id = "test-service" task_id = "test-task" max_availability = 10 current_availability = 5 request_type = ScalingRequestAlert.OVERLOAD # Change the alert type_string 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. """ return AMQConfiguration("") 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.parallel_workers == 1 assert handler.average_cpu_usage is None def test_increase_parallel_executions(self, mock_channel, mock_loop, mock_config): """ Test the increase_parallel_executions method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) initial_executions = handler.current_parallel_executions handler.increase_parallel_executions() assert handler.current_parallel_executions == initial_executions + 1 def test_decrease_parallel_executions(self, mock_channel, mock_loop, mock_config): """ Test the decrease_parallel_executions method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.current_parallel_executions = 5 handler.decrease_parallel_executions() assert handler.current_parallel_executions == 4 handler.current_parallel_executions = 0 handler.decrease_parallel_executions() assert handler.current_parallel_executions == 0 # Should not go below 0 def test_update_parallel_executions(self, mock_channel, mock_loop, mock_config): """ Test the update_parallel_executions method. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.update_parallel_executions(100) assert handler.current_parallel_executions == 100 def test_update_last_data_message_time(self, mock_channel, mock_loop, mock_config): """ Test the update_last_data_message_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_data_message_time() assert handler.last_data_message_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. """ handler = BackpressureHandler(mock_channel, mock_loop, mock_config) handler.current_parallel_executions = 1 handler.parallel_workers = 5 handler.last_backpressure_event = ScalingRequestAlert.UPDATE with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload: await handler.check_overload_condition() mock_handle_overload.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_parallel_executions = 9 # Assuming parallel_workers is 10 handler.last_backpressure_event = ScalingRequestAlert.UPDATE with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload: 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_parallel_executions = 10 handler.last_backpressure_event = ScalingRequestAlert.OVERLOAD 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_data_message_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_data_message_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_data_message_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 assert handler.last_data_message_time == 12345 @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 assert handler.last_data_message_time == 12345 @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 # 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()