amq-adapter-python/tests/adapter/test_backpressure_handler.py

#
# 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


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"
        current_availability = 5

        scale_request = ScaleRequestV1(
            service_id,
            task_id,
            current_availability,
        )

        assert scale_request.version == 1
        assert scale_request.serviceId == service_id
        assert scale_request.taskId == task_id
        assert scale_request.currentAvailability == current_availability
        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"
        current_availability = 5

        scale_request = ScaleRequestV1(
            service_id,
            task_id,
            current_availability,
        )

        json_output = scale_request.to_json()
        expected_json = json.dumps(
            {
                "version": 1,
                "serviceId": service_id,
                "taskId": task_id,
                "currentAvailability": current_availability,
            },
            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"
        current_availability = 5

        scale_request = ScaleRequestV1(
            service_id,
            task_id,
            current_availability,
        )

        json_output = scale_request.to_json()
        expected_json = json.dumps(
            {
                "version": 1,
                "serviceId": service_id,
                "taskId": task_id,
                "currentAvailability": current_availability,
            },
            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 == mock_config.backpressure.threshold_threads

    @pytest.mark.asyncio
    async 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)
        handler.send_backpressure_report = AsyncMock()
        initial_executions = handler.current_parallel_executions

        with patch("time.time") as mock_time:
            mock_time.return_value = 12345
            await handler.increase_parallel_executions()
        assert handler.current_parallel_executions == initial_executions + 1
        assert handler.last_usage_report_time == 12345
        assert handler.send_backpressure_report.call_count == 1

    @pytest.mark.asyncio
    async 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.send_backpressure_report = AsyncMock()
        handler.current_parallel_executions = 5

        with patch("time.time") as mock_time:
            mock_time.return_value = 12345
            await handler.decrease_parallel_executions()
        assert handler.current_parallel_executions == 4
        assert handler.last_recover_report_time == 12345
        assert handler.send_backpressure_report.call_count == 1

        handler.last_recover_report_time = 0
        handler.current_parallel_executions = 0
        await handler.decrease_parallel_executions()
        assert handler.current_parallel_executions == 0  # Should not go below 0

    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_regular_report_loop(self, mock_channel, mock_loop, mock_config):
        """
        Test the _regular_report_loop method.
        """
        handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
        handler.send_backpressure_report = AsyncMock()

        # test 1 loop
        handler.do_loop = 1
        await handler._regular_report_loop()
        handler.send_backpressure_report.assert_called_once()

    @pytest.mark.asyncio
    async def test_send_backpressure_report(self, mock_channel, mock_loop, mock_config):
        """
        Test the send_backpressure_report 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.send_backpressure_report()
        assert handler.publish_backpressure_request.call_count == 1

    @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", 50)
            )
        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", 50)
            )
            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()