Stanislav Hejny
513 lines
20 KiB
Python
513 lines
20 KiB
Python
#
|
|
# 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.
|
|
"""
|
|
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.
|
|
"""
|
|
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()
|