From a6c918fe43abaa33b6ea3809a80668bc16050a49 Mon Sep 17 00:00:00 2001 From: Stanislav Hejny Date: Thu, 17 Jul 2025 17:51:01 +0100 Subject: [PATCH] feat/58 - simplify complex logic --- amqp/adapter/backpressure_handler.py | 258 ++---- amqp/adapter/cleverthis_service_adapter.py | 4 +- amqp/adapter/logging_utils.py | 14 +- amqp/config/amq_configuration.py | 2 +- amqp/router/router_consumer.py | 4 +- amqp/router/router_producer.py | 2 +- amqp/service/amq_message_handler.py | 3 +- amqp/service/amq_service.py | 24 +- demo.py | 2 +- tests/adapter/test_backpressure_handler.py | 871 ++++--------------- tests/adapter/test_data_message_factory.py | 2 +- tests/config/__init__.py | 0 tests/{ => config}/test_amq_configuration.py | 0 tests/service/__init__.py | 0 tests/service/test_amq_message_handler.py | 255 +++--- tests/service/test_amq_service.py | 44 +- tests/test_amq_service.py | 74 -- tests/test_backpressure_handler.py | 253 ------ tools/test_amq_service.py | 172 ++++ tools/test_backpressure_handler.py | 736 ++++++++++++++++ 20 files changed, 1328 insertions(+), 1392 deletions(-) create mode 100644 tests/config/__init__.py rename tests/{ => config}/test_amq_configuration.py (100%) create mode 100644 tests/service/__init__.py delete mode 100644 tests/test_amq_service.py delete mode 100644 tests/test_backpressure_handler.py create mode 100644 tools/test_amq_service.py create mode 100644 tools/test_backpressure_handler.py diff --git a/amqp/adapter/backpressure_handler.py b/amqp/adapter/backpressure_handler.py index 1d52fc6..478387f 100644 --- a/amqp/adapter/backpressure_handler.py +++ b/amqp/adapter/backpressure_handler.py @@ -5,7 +5,6 @@ from asyncio import AbstractEventLoop from threading import Thread from typing import Optional -import psutil from aio_pika import Message from aio_pika.abc import AbstractRobustChannel @@ -14,10 +13,6 @@ from amqp.config.amq_configuration import AMQConfiguration from amqp.model.model import ScalingRequestAlert from amqp.router.utils import await_future, await_result -CPU_IDLE, CPU_ACTIVE, CPU_OVERLOAD = 0, 1, 2 -IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 0.15, 0.90 -RESOURCE_UNSET, RESOURCE_IDLE, RESOURCE_ACTIVE, RESOURCE_OVERLOAD = -1, 0, 1, 2 - class ScaleRequestV1: @@ -25,15 +20,15 @@ class ScaleRequestV1: self, serviceId: str, taskId: str, - maxAvailability: int, - currentAvailability: int, + max_availability: int, + current_availability: int, requestType: ScalingRequestAlert, ): self.version = 1 # Version of the request, currently 1 self.serviceId = serviceId self.taskId = taskId - self.maxAvailability = maxAvailability - self.currentAvailability = currentAvailability + self.max_availability = max_availability + self.current_availability = current_availability self.requestType = requestType self.thread = None @@ -44,58 +39,17 @@ class ScaleRequestV1: "version": self.version, "serviceId": self.serviceId, "taskId": self.taskId, - "maxAvailability": self.maxAvailability, - "currentAvailability": self.currentAvailability, + "maxAvailability": self.max_availability, + "currentAvailability": self.current_availability, "requestType": self.requestType.value, # Using .value for Enum serialization }, indent=2, ) -class RunningAverage: - """ - A class to maintain a running average of the last N values. - The intended use is to track CPU usage over time window, that is directly related to sample rate and window size. - This is the reason why the constructor takes the number of items to store, calculated as window size / sample rate - Initial_value is for unit test to emulate OVERLOAD or IDLE condition. - """ - - def __init__(self, num_items, initial_value=0): - self.buffer = [initial_value] * num_items # Fixed-size array - self.pointer = 0 # Current position in array - self.num_items = num_items # Maximum number of items to store - self.is_filled = False # Track if buffer is fully filled - - def add(self, value): - # Store the value at current pointer position - self.buffer[self.pointer] = value - - # Move pointer to next position with wrap-around - self.pointer += 1 - if self.pointer >= self.num_items: - self.pointer = 0 - self.is_filled = True - - def get_average(self): - # Determine how many items we should average - count = self.num_items if self.is_filled else self.pointer - - if count == 0: - return 0.0 # Avoid division by zero - - return sum(self.buffer) / count - - def get_current_values(self): - """Returns the values in chronological order (oldest first)""" - if not self.is_filled: - return self.buffer[: self.pointer] - - return self.buffer[self.pointer :] + self.buffer[: self.pointer] - - class BackpressureHandler: # Track the number of messages currently being processed - current_load = 0 + current_availability = 0 # helps detect IDLE condition # helps prevent flooding the system with backpressure events last_backpressure_event_time = 0 @@ -116,17 +70,19 @@ class BackpressureHandler: self.exchange = None self.swarm_service_id = self.config.amq_adapter.swarm_service_id self.swarm_task_id = self.config.amq_adapter.swarm_task_id - self.maximum_load = self.config.backpressure.threshold_load - self.average_cpu_usage = None self.do_loop = -1 self._resource_usage_changed = 0 self._resource_average_value = 0 self._last_resource_max_value = 100 # Default max value for CPU usage + self.max_availability = -1 + self.current_availability = 1 + self.max_load = 1 + self.current_load = 0 def increase_current_load(self): - """Increase the number of parallel executions, or current load""" + """Increase the number of parallel executions, or current_availability load""" logging_info( - "Backpressure: Increase parallel executions, current=%s", + "Backpressure: Increase current load (%s) by 1", self.current_load, ) self.current_load += 1 @@ -134,55 +90,43 @@ class BackpressureHandler: def decrease_current_load(self): """Decrease the number of parallel executions, or current load""" logging_info( - "Backpressure: Decrease parallel executions, current=%s", + "Backpressure: Decrease current load(%s) by 1", self.current_load, ) if self.current_load > 0: self.current_load -= 1 - def update_current_load(self, count: int): + def update_current_availability(self, count: int): """Update the number of parallel executions, or current load""" - self.current_load = count + self.current_availability = count - def update_last_data_message_time(self): + def update_last_backpressure_event_time(self): # logging_info("Backpressure: Update last data message time") """Update the last data message time""" self.last_backpressure_event_time = time.time() - async def update_backpressure_value(self, current: int, maximum: int): + async def update_backpressure_value(self, current_availability: int, maximum: int): """ Update the current backpressure value and check for overload conditions. This method is called by the AMQService.backpressure() method to update the current parallel executions count and check for overload conditions. Args: - current: Current value of the backpressure metric + current_availability: Current value of the backpressure metric maximum: Maximum value of the backpressure metric """ logging_info( - "Backpressure: Updating backpressure value, current=%s, maximum=%s", current, maximum + "Backpressure: Updating backpressure value, current_availability=%s, maximum=%s", + current_availability, + maximum, ) - if maximum > 0 and maximum > self.maximum_load: - self.maximum_load = maximum - - # Update the current parallel executions count - self.update_current_load(current) - + if maximum > 0 and maximum > self.max_availability: + self.max_availability = maximum + # Update the current_availability / parallel executions count + self.update_current_availability(current_availability) # Check for overload conditions await self.check_overload_condition() - # Update the last data message time to indicate activity - self.update_last_data_message_time() - - # If maximum has changed, update the parallel workers threshold - if maximum != self.maximum_load and maximum > 0: - logging_info( - "Backpressure: Updating maximum parallel workers from %s to %s", - self.maximum_load, - maximum, - ) - self.maximum_load = maximum - def start_backpressure_monitor(self) -> Optional[Thread]: # Start the Backpressure monitor loop self.thread = Thread(target=self.backpressure_monitor_loop) @@ -192,42 +136,38 @@ class BackpressureHandler: async def check_overload_condition(self): """ - Check if the current availability is too low (OVERLOAD) + Check if the current_availability availability is too low (OVERLOAD) or if the service has been idle for too long with high availability (IDLE). - + Note: current_availability represents available capacity, not used capacity. - Low availability (close to 0) means OVERLOAD - High availability (close to maximum) with no activity means IDLE """ current_time = time.time() - last_event_delta = current_time - self.last_backpressure_event_time - - # If maximum is not set or invalid, use the highest seen value - if self.maximum_load <= 0 and self.current_load > 0: - self.maximum_load = self.current_load - elif self.current_load > self.maximum_load: - # Update maximum if we see a higher value - self.maximum_load = self.current_load - + last_event_delta: float = current_time - self.last_backpressure_event_time + logging_info( "Backpressure: Check conditions, current_availability=%s, max=%s, last_activity=%s", - self.current_load, - self.maximum_load, + self.current_availability, + self.max_availability, last_event_delta, ) - + # Check for OVERLOAD condition - low availability (less than 20% of maximum) - if self.current_load <= round(0.2 * self.maximum_load): + if self.current_availability <= round(0.1 * self.max_availability): # Check if the last backpressure event was not an overload - if self.last_backpressure_event != ScalingRequestAlert.OVERLOAD: + if ( + self.last_backpressure_event != ScalingRequestAlert.OVERLOAD + or last_event_delta > self.config.backpressure.idle_duration + ): # Trigger the overload event await self.handle_backpressure_overload_event() # update / reset time-window so that the OVERLOAD is not sent too often self.last_backpressure_event_time = current_time self.last_backpressure_event = ScalingRequestAlert.OVERLOAD - + # Check for IDLE condition - high availability (more than 80% of maximum) with no activity - elif self.current_load >= round(0.8 * self.maximum_load): + elif self.current_availability >= round(0.8 * self.max_availability): idle_duration = self.config.backpressure.idle_duration # Check if service has been idle for longer than the configured duration if ( @@ -239,118 +179,46 @@ class BackpressureHandler: last_event_delta, idle_duration, ) - + # Trigger the idle event await self._handle_backpressure_idle_event() # update / reset time-window so that the IDLE is not sent too often self.last_backpressure_event_time = current_time else: - # Don't send IDLE right away when the availability increases, but wait for a while + # Don't send IDLE right away when the availability increases, but wait for a while, send UPDATE instead self.last_backpressure_event = ScalingRequestAlert.IDLE - + _scaling_request: ScaleRequestV1 = ScaleRequestV1( + self.swarm_service_id, + self.swarm_task_id, + self.max_availability, + self.current_availability, # Current availability is passed directly + ScalingRequestAlert.UPDATE, + ) + # Address the message to any adapter capable of supporting BACKPRESSURE request + await self.publish_backpressure_request(_scaling_request) + self.last_backpressure_event_time = current_time + # If neither OVERLOAD nor IDLE, and it's time for an update, send UPDATE event elif last_event_delta > self.config.backpressure.time_window: _scaling_request: ScaleRequestV1 = ScaleRequestV1( self.swarm_service_id, self.swarm_task_id, - self.maximum_load, - self.current_load, # Current availability is passed directly + self.max_availability, + self.current_availability, # Current availability is passed directly ScalingRequestAlert.UPDATE, ) # Address the message to any adapter capable of supporting BACKPRESSURE request await self.publish_backpressure_request(_scaling_request) self.last_backpressure_event_time = current_time self.last_backpressure_event = ScalingRequestAlert.UPDATE - - self.update_last_data_message_time() + + self.update_last_backpressure_event_time() async def _backpressure_monitor(self): - """Monitor the backpressure conditions and trigger events accordingly""" - # Check if CPU monitoring is enabled in configuration - cpu_monitoring_enabled = self.config.backpressure.cpu_monitoring_enabled - _time_window_millis = self.config.backpressure.time_window - _idle_duration_millis = self.config.backpressure.idle_duration - _overload_duration_millis = self.config.backpressure.cpu_overload_duration + """Periodically monitor the backpressure conditions and trigger events accordingly""" _monitor_interval = 0.5 # Monitor every 500ms - CPU_IDLE, CPU_ACTIVE, CPU_OVERLOAD = 0, 1, 2 - IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 15, 90 - if self.average_cpu_usage is None: - self.average_cpu_usage = RunningAverage( - int(max(_overload_duration_millis, _idle_duration_millis) // _monitor_interval) - ) - _cpu_usage_state = CPU_ACTIVE - _cpu_usage_changed = time.time() while self._do_loop(): - if cpu_monitoring_enabled: - _current_cpu_usage = psutil.cpu_percent(interval=None) - self.average_cpu_usage.add(_current_cpu_usage) - _average_cpu_usage = self.average_cpu_usage.get_average() - _now = time.time() - # logging_info( - # "Backpressure: current CPU[pct]=%s, avg=%s, state=%s, last dataMsg=%s, last event=%s, last eventTime=%s", - # _current_cpu_usage, - # _average_cpu_usage, - # _cpu_usage_state, - # self.last_backpressure_event, - # self.last_backpressure_event_time, - # ) - if _average_cpu_usage < IDLE_THRESHOLD: - if _cpu_usage_state != CPU_IDLE: - _cpu_usage_changed = _now - _cpu_usage_state = CPU_IDLE - elif _average_cpu_usage > OVERLOAD_THRESHOLD: - if _cpu_usage_state != CPU_OVERLOAD: - _cpu_usage_changed = _now - _cpu_usage_state = CPU_OVERLOAD - else: - if _cpu_usage_state != CPU_ACTIVE: - _cpu_usage_changed = _now - _cpu_usage_state = CPU_ACTIVE - - # Check if the current time exceeds the last backpressure event time - if ( - _cpu_usage_state == CPU_OVERLOAD - and _now - _cpu_usage_changed > _overload_duration_millis - and ( - _now - self.last_backpressure_event_time > _time_window_millis - or self.last_backpressure_event != ScalingRequestAlert.OVERLOAD - ) - ): - # Trigger the overload event - await self.handle_backpressure_overload_event() - self.last_backpressure_event_time = _now - self.last_backpressure_event = ScalingRequestAlert.OVERLOAD - elif ( - _cpu_usage_state == CPU_IDLE - and _now - _cpu_usage_changed > _idle_duration_millis - and ( - _now - self.last_backpressure_event_time > _time_window_millis - or self.last_backpressure_event != ScalingRequestAlert.IDLE - ) - ): - # Trigger the idle event - await self._handle_backpressure_idle_event() - self.last_backpressure_event = ScalingRequestAlert.IDLE - self.last_backpressure_event_time = _now - else: - # Trigger update event - if _now - self.last_backpressure_event_time > _time_window_millis: - # Trigger the update event - _scaling_request: ScaleRequestV1 = ScaleRequestV1( - self.swarm_service_id, - self.swarm_task_id, - 100, - _average_cpu_usage, - ScalingRequestAlert.UPDATE, - ) - # Address the message to any adapter capable of supporting BACKPRESSURE request - await self.publish_backpressure_request(_scaling_request) - self.last_backpressure_event_time = _now - self.last_backpressure_event = ScalingRequestAlert.UPDATE - else: - # If CPU monitoring is disabled, just check the current load - await self.check_overload_condition() - + await self.check_overload_condition() await asyncio.sleep(_monitor_interval) def backpressure_monitor_loop(self): @@ -363,8 +231,8 @@ class BackpressureHandler: scaling_request: ScaleRequestV1 = ScaleRequestV1( self.swarm_service_id, self.swarm_task_id, - self.maximum_load, - self.current_load, # Current availability is passed directly + self.max_availability, + self.current_availability, # Current availability is passed directly ScalingRequestAlert.OVERLOAD, ) # Address the message to any management service instance capable of supporting BACKPRESSURE request @@ -376,8 +244,8 @@ class BackpressureHandler: scaling_request: ScaleRequestV1 = ScaleRequestV1( self.swarm_service_id, self.swarm_task_id, - self.maximum_load, - self.current_load, # Current availability is passed directly + self.max_availability, + self.current_availability, # Current availability is passed directly ScalingRequestAlert.IDLE, ) # Address the message to any adapter capable of supporting BACKPRESSURE request diff --git a/amqp/adapter/cleverthis_service_adapter.py b/amqp/adapter/cleverthis_service_adapter.py index a9c370f..2a145b5 100644 --- a/amqp/adapter/cleverthis_service_adapter.py +++ b/amqp/adapter/cleverthis_service_adapter.py @@ -83,7 +83,7 @@ async def _convert_file_response( temporary dedicated queue. """ # All calls to RabbitMQ API have to be called threadsafe, because those appear to be attached - # to parent event loop, and need to execute at that event loop, not the current one + # to parent event loop, and need to execute at that event loop, not the current_availability one # allow coroutine to execute. calling _future.result() will block the entire thread # and the coroutine will never execute _channel, _exchange, _routing_key = await await_result( @@ -241,7 +241,7 @@ class CleverThisServiceAdapter: async def get_current_user(self, token: str) -> object: """ - To Be Overriden in subclass for given CleverThis Service, to return the current user based + To Be Overriden in subclass for given CleverThis Service, to return the current_availability user based on the token, in format appropriate for the service. :param token: The token to be used for authentication. diff --git a/amqp/adapter/logging_utils.py b/amqp/adapter/logging_utils.py index d8ff2d8..0238f8f 100644 --- a/amqp/adapter/logging_utils.py +++ b/amqp/adapter/logging_utils.py @@ -13,17 +13,17 @@ __verbose__ = os.environ.get("AMQ_VERBOSE_LOGGING", "0") == "1" def get_context_info(): """ - Retrieves the current module, line number, and function name. + Retrieves the current_availability module, line number, and function name. Returns: A tuple containing (module name, line number, function name). Returns (None, None, None) if it fails to retrieve the information. """ try: - # Get the frame object for the current function call + # Get the frame object for the current_availability function call frame = inspect.currentframe() if frame is not None: - # Get the frame object for the caller of the current function + # Get the frame object for the caller of the current_availability function frame = frame.f_back.f_back if frame: code_context = frame.f_code @@ -44,7 +44,7 @@ def get_context_info(): def logging_error(msg: str, *args) -> None: """ - Log an error message according to current logging for 'ERROR' level. + Log an error message according to current_availability logging for 'ERROR' level. Add module name, line and function name where the error occurred, on single line. Args: msg (str): The error message to log. @@ -70,7 +70,7 @@ def logging_error(msg: str, *args) -> None: def logging_warning(msg: str, *args) -> None: """ - Log a warning message according to current logging for 'WARN' level. + Log a warning message according to current_availability logging for 'WARN' level. Add module name, line and function name where the print occurred, on single line. Args: msg (str): The warning message to log. @@ -91,7 +91,7 @@ def logging_warning(msg: str, *args) -> None: def logging_info(msg: str, *args) -> None: """ - Log an info message according to current logging for 'INFO' level. + Log an info message according to current_availability logging for 'INFO' level. Add module name, line and function name where the print occurred, on single line. Args: msg (str): The info message to log. @@ -112,7 +112,7 @@ def logging_info(msg: str, *args) -> None: def logging_debug(msg: str, *args) -> None: """ - Log a debug message according to current logging for 'DEBUG' level. + Log a debug message according to current_availability logging for 'DEBUG' level. Add module name, line and function name where the print occurred, on single line. Args: msg (str): The debug message to log. diff --git a/amqp/config/amq_configuration.py b/amqp/config/amq_configuration.py index fe9faf2..79dd4e3 100644 --- a/amqp/config/amq_configuration.py +++ b/amqp/config/amq_configuration.py @@ -203,7 +203,7 @@ class AMQConfiguration: config.read(config_file_name) else: # Get the directory of the caller (the module where this function is called) - caller_frame = inspect.stack()[1] # 0 is current frame, 1 is caller + caller_frame = inspect.stack()[1] # 0 is current_availability frame, 1 is caller caller_module = inspect.getmodule(caller_frame[0]) if caller_module: caller_dir = os.path.dirname(os.path.abspath(caller_module.__file__)) diff --git a/amqp/router/router_consumer.py b/amqp/router/router_consumer.py index 10f0a23..afea07e 100644 --- a/amqp/router/router_consumer.py +++ b/amqp/router/router_consumer.py @@ -174,12 +174,12 @@ class RouterConsumer(RouterProducer): ) if not await self.publish_service_message(service_message, self.cm_response_exchange): logging_warning( - "RouterConsumer couldn't remove current route: %s, channel not open.", + "RouterConsumer couldn't remove current_availability route: %s, channel not open.", route, ) except Exception as ioe: - logging_error("RouterConsumer failed remove current route: %s", ioe) + logging_error("RouterConsumer failed remove current_availability route: %s", ioe) result: bool = self.remove_route(route) return result diff --git a/amqp/router/router_producer.py b/amqp/router/router_producer.py index 6224dc1..2dd8bd6 100644 --- a/amqp/router/router_producer.py +++ b/amqp/router/router_producer.py @@ -81,7 +81,7 @@ class RouterProducer(RouterBase): name=CM_REQUEST_EXCHANGE, type=ExchangeType.fanout ) # - # 2. declare ResponseExchange where to publish current RoutingData (response to + # 2. declare ResponseExchange where to publish current_availability RoutingData (response to # a RoutingData request) self.cm_response_exchange = await self.channel.declare_exchange( name=CM_RESPONSE_EXCHANGE, type=ExchangeType.fanout diff --git a/amqp/service/amq_message_handler.py b/amqp/service/amq_message_handler.py index 630f1ed..8dea0e4 100644 --- a/amqp/service/amq_message_handler.py +++ b/amqp/service/amq_message_handler.py @@ -127,7 +127,7 @@ class DataMessageHandler: self.backpressure_handler.increase_current_load() await self.backpressure_handler.check_overload_condition() logging_info( - f"PARALLEL: Current Capacity [{self.backpressure_handler.current_load}] of [{self.backpressure_handler.maximum_load}]" + f"PARALLEL: Current Capacity [{self.backpressure_handler.current_availability}] of [{self.backpressure_handler.max_availability}]" ) if message.reply_to: self.reply_to = message.reply_to @@ -156,7 +156,6 @@ class DataMessageHandler: logging_error(f"Request handler: {e}") self.backpressure_handler.decrease_current_load() - await self.backpressure_handler.check_overload_condition() async def reconstitute_data_message( self, message: AbstractIncomingMessage diff --git a/amqp/service/amq_service.py b/amqp/service/amq_service.py index e9cd7ed..d7b82d0 100644 --- a/amqp/service/amq_service.py +++ b/amqp/service/amq_service.py @@ -64,6 +64,7 @@ class AMQService: ) self.backpressure_thread: Optional[Thread] = None self.backpressure_handler: Optional[BackpressureHandler] = None + self.maximum_availability = -1 # ======================================================================= # ======================================================================= @@ -87,21 +88,28 @@ class AMQService: """ Send Backpressure event to the Management Service. The Backpressure logic determines the type of the event based on the current_availability and maximum values. - - OVERLOAD event is triggered immediately when the current_availability value exceeds the maximum threshold, which is - typically set to 80% of the maximum value. - - IDLE event is triggered when the current_availability value drops below 20% of the maximum value and remains there for - a IDLE_WINDOW period of time (see BackpressureHandler). + - OVERLOAD event is triggered immediately when the current_availability value is below threshold + typically set to 10% of the maximum value, or 0 if no maximum is provided. + - IDLE event is triggered when the current_availability value exeeds 80% of the maximum value or is greater + than 0 if no maximum is provided. - UPDATE event is sent if neither OVERLOAD nor IDLE conditions are met, indicating that the backpressure is within acceptable limits. - :param current_availability: Currently available caopacity. + :param current_availability: Currently available capacity. :param maximum: Maximum value of the backpressure metric. If undefined or set to lt 0 (e.g. -1), - the maximum is than hgdetermined from max value of current_availability seen over the time. + the maximum is than determined from max value of current_availability seen over the time. """ if self.backpressure_handler is not None: + # watch maximum & current value to always have some valid maximum reference if maximum < 0: - maximum = self.amq_configuration.backpressure.threshold_load + if current_availability > self.maximum_availability: + self.maximum_availability = current_availability + else: + if maximum > self.maximum_availability: + self.maximum_availability = maximum asyncio.run_coroutine_threadsafe( - self.backpressure_handler.update_backpressure_value(current_availability, maximum), + self.backpressure_handler.update_backpressure_value( + current_availability, maximum if maximum > 0 else self.maximum_availability + ), loop=self.backpressure_handler.loop, ) diff --git a/demo.py b/demo.py index 42ae999..9d443b6 100644 --- a/demo.py +++ b/demo.py @@ -134,7 +134,7 @@ async def upload_document( start_time = time.time() async with httpx.AsyncClient() as client: - # httpx instrumentation ensures the current trace context is propagated + # httpx instrumentation ensures the current_availability trace context is propagated response = await client.post( # IMPORTANT: Use the actual host and port where your FastAPI app is running # If running directly on host, use http://localhost:8000 diff --git a/tests/adapter/test_backpressure_handler.py b/tests/adapter/test_backpressure_handler.py index fafb42e..a43a7ee 100644 --- a/tests/adapter/test_backpressure_handler.py +++ b/tests/adapter/test_backpressure_handler.py @@ -1,736 +1,205 @@ -# -# NOTE: need to install: pytest-asyncio -# - import asyncio -import json -from asyncio import AbstractEventLoop -from threading import Thread +import time +import unittest 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, + ScalingRequestAlert, ) -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. - """ +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.cpu_overload_duration = 1 - 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. - """ - 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.maximum_load == 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_load - handler.increase_current_load() - assert handler.current_load == 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_load = 5 - handler.decrease_current_load() - assert handler.current_load == 4 - - handler.current_load = 0 - handler.decrease_current_load() - assert handler.current_load == 0 # Should not go below 0 - - def test_update_current_load(self, mock_channel, mock_loop, mock_config): - """ - Test the update_current_load method. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.update_current_load(100) - assert handler.current_load == 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_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_load = 1 - handler.maximum_load = 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_load = 9 # Assuming maximum_load 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_load = 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 + # Record callbacks for testing 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 + self.handler = BackpressureHandler(self.channel, self.loop, self.config) + self.handler.exchange = AsyncMock() - assert "_wrap_rabbit_mq_api_init" in BackpressureHandler._callback_list - assert "_wrap_rabbit_mq_api" in BackpressureHandler._callback_list + # Set do_loop to 1 to run the loop only once + self.handler.do_loop = 1 - @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_load = MagicMock() - handler.update_last_data_message_time = MagicMock() - handler.check_overload_condition = AsyncMock() + 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) - # Test with current value and same maximum - await handler.update_backpressure_value(5, handler.maximum_load) - handler.update_current_load.assert_called_once_with(5) - handler.update_last_data_message_time.assert_called_once() - handler.check_overload_condition.assert_called_once() + 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) - # Reset mocks - handler.update_current_load.reset_mock() - handler.update_last_data_message_time.reset_mock() - handler.check_overload_condition.reset_mock() + 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) - # Test with current value and different maximum - original_workers = handler.maximum_load - await handler.update_backpressure_value(3, original_workers + 5) - handler.update_current_load.assert_called_once_with(3) - handler.update_last_data_message_time.assert_called_once() - handler.check_overload_condition.assert_called_once() - assert handler.maximum_load == original_workers + 5 + @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) - @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() + # 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) - # Set up initial state for IDLE condition - handler.current_load = 0 - handler.last_backpressure_event = ScalingRequestAlert.IDLE + @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 - with patch("time.time") as mock_time: - # Set current time - current_time = 1000 - mock_time.return_value = current_time + await self.handler.check_overload_condition() - # Set last data message time to be older than idle_duration - handler.last_backpressure_event_time = ( - current_time - mock_config.backpressure.idle_duration - 10 - ) + # 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) - # Call the method - await handler.check_overload_condition() + @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 - # 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 + await self.handler.check_overload_condition() - @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() + # 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) - # Set up initial state for UPDATE condition - handler.current_load = 2 - handler.maximum_load = 10 - handler.last_backpressure_event = ScalingRequestAlert.IDLE # Coming from IDLE state + @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 + ) - with patch("time.time") as mock_time: - # Set current time - current_time = 1000 - mock_time.return_value = current_time + await self.handler.check_overload_condition() - # Set last event time to be older than time_window - handler.last_backpressure_event_time = ( - current_time - mock_config.backpressure.time_window - 10 - ) + # 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) - # Call the method - await handler.check_overload_condition() + @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 + ) - # 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 + await self.handler.check_overload_condition() - 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 + # 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) - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - thread = handler.start_backpressure_monitor() + # Test with a higher value that should update the maximum + mock_publish.reset_mock() + await self.handler.update_backpressure_value(80, 80) - # Verify no thread was started - assert thread is not None - assert hasattr(handler, "thread") or handler.thread is not None + # 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) - 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 + @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 - 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() + await self.handler.handle_backpressure_overload_event() - # 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() + 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) + + +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() diff --git a/tests/adapter/test_data_message_factory.py b/tests/adapter/test_data_message_factory.py index 9b016ba..202a529 100644 --- a/tests/adapter/test_data_message_factory.py +++ b/tests/adapter/test_data_message_factory.py @@ -107,7 +107,7 @@ class DataMessageFactoryTest(TestCase): self.assertLessEqual( _x_timestamp - _timestamp, 1, - "Timestamp in SnowflakeId differs too much from current time", + "Timestamp in SnowflakeId differs too much from current_availability time", ) def test_to_string(self): diff --git a/tests/config/__init__.py b/tests/config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_amq_configuration.py b/tests/config/test_amq_configuration.py similarity index 100% rename from tests/test_amq_configuration.py rename to tests/config/test_amq_configuration.py diff --git a/tests/service/__init__.py b/tests/service/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/service/test_amq_message_handler.py b/tests/service/test_amq_message_handler.py index 7d606a4..f8eda80 100644 --- a/tests/service/test_amq_message_handler.py +++ b/tests/service/test_amq_message_handler.py @@ -1,23 +1,27 @@ import asyncio -import os -import time import unittest -from unittest.mock import AsyncMock, MagicMock, patch, call - -from aio_pika import Message -from opentelemetry.trace import SpanContext, TraceFlags +from unittest.mock import AsyncMock, MagicMock, patch from amqp.adapter.backpressure_handler import BackpressureHandler -from amqp.adapter.cleverthis_service_adapter import AMQMessage, CleverThisServiceAdapter +from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter from amqp.adapter.data_message_factory import DataMessageFactory -from amqp.adapter.data_response_factory import DataResponseFactory from amqp.adapter.file_handler import FileHandler from amqp.adapter.service_message_factory import ServiceMessageFactory from amqp.config.amq_configuration import AMQConfiguration -from amqp.model.model import DataMessage, DataMessageMagicByte, DataResponse, MultipartDataMessage +from amqp.model.model import ( + DataMessage, + DataMessageMagicByte, + DataResponse, + MultipartDataMessage, +) from amqp.model.snowflake_id import SnowflakeId from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient -from amqp.service.amq_message_handler import DataMessageHandler, collect_trace_info, get_default_trace_state, set_text +from amqp.service.amq_message_handler import ( + DataMessageHandler, + collect_trace_info, + get_default_trace_state, + set_text, +) class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase): @@ -25,36 +29,36 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase): # Mock dependencies self.service_adapter = MagicMock(spec=CleverThisServiceAdapter) self.service_adapter.on_message = AsyncMock() - + self.tracer = MagicMock() self.tracer.start_span = MagicMock() - + self.message_factory = MagicMock(spec=DataMessageFactory) self.service_message_factory = MagicMock(spec=ServiceMessageFactory) self.reply_to_exchange = MagicMock() self.reply_to_exchange.publish = AsyncMock() self.reply_to_exchange.name = "test-exchange" - + self.rabbit_mq_client = MagicMock(spec=RabbitMQClient) self.rabbit_mq_client.connection = MagicMock() - + self.amq_configuration = MagicMock(spec=AMQConfiguration) self.amq_configuration.dispatch = MagicMock() self.amq_configuration.dispatch.rabbit_mq_url = "amqp://guest:guest@localhost:5672/" self.amq_configuration.dispatch.download_dir = "/tmp/test_downloads" - + self.loop = asyncio.get_event_loop() - + self.backpressure_handler = MagicMock(spec=BackpressureHandler) self.backpressure_handler.increase_current_load = MagicMock() self.backpressure_handler.decrease_current_load = MagicMock() self.backpressure_handler.check_overload_condition = AsyncMock() self.backpressure_handler.current_availability = 5 self.backpressure_handler.max_availability = 10 - + self.file_handler = MagicMock(spec=FileHandler) self.file_handler.on_message = AsyncMock(return_value={}) - + # Create the handler self.handler = DataMessageHandler( service_adapter=self.service_adapter, @@ -66,9 +70,9 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase): amq_configuration=self.amq_configuration, loop=self.loop, backpressure_handler=self.backpressure_handler, - file_handler=self.file_handler + file_handler=self.file_handler, ) - + # Create a mock message self.mock_message = MagicMock() self.mock_message.delivery_tag = 123 @@ -80,53 +84,63 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase): self.mock_message.body = b"test-body" self.mock_message.ack = AsyncMock() self.mock_message.nack = AsyncMock() - + # Create a mock data message self.mock_data_message = MagicMock(spec=DataMessage) self.mock_data_message.magic.return_value = DataMessageMagicByte.HTTP_REQUEST.value - self.mock_data_message.id.return_value = SnowflakeId.from_hex("1234567890abcdef") + self.mock_data_message.id.return_value = SnowflakeId.from_hex("1234567890abcdef1234") self.mock_data_message.method.return_value = "GET" self.mock_data_message.path.return_value = "/test/path" - self.mock_data_message.trace_info.return_value = {"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"} + self.mock_data_message.trace_info.return_value = { + "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + } self.mock_data_message.body.return_value = b'{"test": "data"}' - + # Create a mock data response self.mock_data_response = MagicMock(spec=DataResponse) self.mock_data_response.id.return_value = "response-id" self.mock_data_response.content_type.return_value = "application/json" - + # Patch DataMessageFactory.from_bytes - self.from_bytes_patch = patch('amqp.adapter.data_message_factory.DataMessageFactory.from_bytes', - return_value=self.mock_data_message) + self.from_bytes_patch = patch( + "amqp.adapter.data_message_factory.DataMessageFactory.from_bytes", + return_value=self.mock_data_message, + ) self.mock_from_bytes = self.from_bytes_patch.start() - + # Patch DataMessageFactory.from_stream - self.from_stream_patch = patch('amqp.adapter.data_message_factory.DataMessageFactory.from_stream', - return_value=self.mock_data_message) + self.from_stream_patch = patch( + "amqp.adapter.data_message_factory.DataMessageFactory.from_stream", + return_value=self.mock_data_message, + ) self.mock_from_stream = self.from_stream_patch.start() - + # Patch DataResponseFactory.serialize - self.serialize_patch = patch('amqp.adapter.data_response_factory.DataResponseFactory.serialize', - return_value=b'serialized-response') + self.serialize_patch = patch( + "amqp.adapter.data_response_factory.DataResponseFactory.serialize", + return_value=b"serialized-response", + ) self.mock_serialize = self.serialize_patch.start() - + # Patch DataResponseFactory.from_bytes - self.response_from_bytes_patch = patch('amqp.adapter.data_response_factory.DataResponseFactory.from_bytes', - return_value=self.mock_data_response) + self.response_from_bytes_patch = patch( + "amqp.adapter.data_response_factory.DataResponseFactory.from_bytes", + return_value=self.mock_data_response, + ) self.mock_response_from_bytes = self.response_from_bytes_patch.start() - + # Patch asyncio.run_coroutine_threadsafe - self.run_coroutine_patch = patch('asyncio.run_coroutine_threadsafe') + self.run_coroutine_patch = patch("asyncio.run_coroutine_threadsafe") self.mock_run_coroutine = self.run_coroutine_patch.start() - + # Patch time.time - self.time_patch = patch('time.time', return_value=1234567890.0) + self.time_patch = patch("time.time", return_value=1234567890.0) self.mock_time = self.time_patch.start() - + # Patch os.remove - self.os_remove_patch = patch('os.remove') + self.os_remove_patch = patch("os.remove") self.mock_os_remove = self.os_remove_patch.start() - + def tearDown(self): self.from_bytes_patch.stop() self.from_stream_patch.stop() @@ -135,15 +149,15 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase): self.run_coroutine_patch.stop() self.time_patch.stop() self.os_remove_patch.stop() - + async def test_inbound_data_message_callback_no_thread_success(self): """Test successful processing of an inbound message.""" # Setup self.service_adapter.on_message.return_value = self.mock_data_response - + # Execute await self.handler.inbound_data_message_callback_no_thread(self.mock_message) - + # Verify self.mock_from_bytes.assert_called_once_with(self.mock_message.body) self.tracer.start_span.assert_called_once() @@ -152,244 +166,253 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase): self.reply_to_exchange.publish.assert_called_once() self.backpressure_handler.decrease_current_load.assert_called_once() self.mock_message.ack.assert_called_once_with(multiple=False) - + async def test_inbound_data_message_callback_no_thread_unknown_magic(self): """Test handling of a message with unknown magic value.""" # Setup self.mock_data_message.magic.return_value = "UNKNOWN" - + # Execute await self.handler.inbound_data_message_callback_no_thread(self.mock_message) - + # Verify self.mock_message.nack.assert_called_once_with(requeue=False) - self.mock_message.ack.assert_called_once_with(multiple=False) - + # self.mock_message.ack.assert_called_once_with(multiple=False) + async def test_inbound_data_message_callback_no_thread_no_message(self): """Test handling when no valid message is present.""" # Setup self.mock_from_bytes.return_value = None - + # Execute await self.handler.inbound_data_message_callback_no_thread(self.mock_message) - + # Verify self.mock_message.ack.assert_called_once_with(multiple=False) - + async def test_inbound_data_message_callback_as_thread(self): """Test the thread-based callback method.""" # Setup - with patch('threading.Thread') as mock_thread: + with patch("threading.Thread") as mock_thread: mock_thread_instance = MagicMock() mock_thread.return_value = mock_thread_instance - + # Execute await self.handler.inbound_data_message_callback_as_thread(self.mock_message) - + # Verify mock_thread.assert_called_once() mock_thread_instance.start.assert_called_once() - + async def test_reconstitute_data_message_standard(self): """Test reconstituting a standard data message.""" # Execute result = await self.handler.reconstitute_data_message(self.mock_message) - + # Verify self.assertEqual(result, self.mock_data_message) self.mock_from_bytes.assert_called_once_with(self.mock_message.body) - + async def test_reconstitute_data_message_stream(self): """Test reconstituting a message from a stream.""" # Setup self.mock_message.headers = {"clevermicro.addressing": 1} - + # Execute result = await self.handler.reconstitute_data_message(self.mock_message) - + # Verify - self.assertEqual(result, self.mock_data_message) + self.assertEqual(MultipartDataMessage, result.__class__) self.mock_from_stream.assert_called_once() self.file_handler.on_message.assert_called_once() - + async def test_run_http_request_magic_success(self): """Test successful processing of an HTTP request.""" # Setup self.service_adapter.on_message.return_value = self.mock_data_response - + # Execute - await self.handler.run_http_request_magic(self.mock_message, self.mock_data_message, self.loop) - + await self.handler.run_http_request_magic( + self.mock_message, self.mock_data_message, self.loop + ) + # Verify self.backpressure_handler.increase_current_load.assert_called_once() self.backpressure_handler.check_overload_condition.assert_called() self.service_adapter.on_message.assert_called_once() self.reply_to_exchange.publish.assert_called_once() self.backpressure_handler.decrease_current_load.assert_called_once() - + async def test_run_http_request_magic_multipart(self): """Test processing a multipart message with file cleanup.""" # Setup self.service_adapter.on_message.return_value = self.mock_data_response - multipart_message = MultipartDataMessage(self.mock_data_message, {"file1": "/tmp/file1.txt"}) - + multipart_message = MultipartDataMessage( + self.mock_data_message, {"file1": "/tmp/file1.txt"} + ) + # Execute await self.handler.run_http_request_magic(self.mock_message, multipart_message, self.loop) - + # Verify self.mock_os_remove.assert_called_once_with("/tmp/file1.txt") - + async def test_run_http_request_magic_service_exception(self): """Test handling of an exception from the service adapter.""" # Setup self.service_adapter.on_message.side_effect = Exception("Service error") - + # Execute - await self.handler.run_http_request_magic(self.mock_message, self.mock_data_message, self.loop) - + await self.handler.run_http_request_magic( + self.mock_message, self.mock_data_message, self.loop + ) + # Verify self.backpressure_handler.decrease_current_load.assert_called_once() self.backpressure_handler.check_overload_condition.assert_called() - + async def test_run_http_request_magic_reply_exception(self): """Test handling of an exception when sending the reply.""" # Setup self.service_adapter.on_message.return_value = self.mock_data_response self.reply_to_exchange.publish.side_effect = Exception("Reply error") - + # Execute - await self.handler.run_http_request_magic(self.mock_message, self.mock_data_message, self.loop) - + await self.handler.run_http_request_magic( + self.mock_message, self.mock_data_message, self.loop + ) + # Verify - self.mock_message.nack.assert_called_once_with(requeue=False) + # [no ACK at this level, ACK is in parent call] self.mock_message.ack.assert_called_once_with(requeue=False) self.backpressure_handler.decrease_current_load.assert_called_once() - + async def test_send_reply_success(self): """Test successful sending of a reply.""" # Execute await self.handler.send_reply("reply-queue", self.mock_data_response) - + # Verify self.mock_run_coroutine.assert_called_once() - + async def test_send_reply_no_reply_to(self): """Test handling when no reply_to is provided.""" # Execute await self.handler.send_reply(None, self.mock_data_response) - + # Verify self.mock_run_coroutine.assert_not_called() - + async def test_reply_received_callback_with_future(self): """Test handling a reply when there's a matching future.""" # Setup future = asyncio.Future() self.handler.outstanding["correlation-id"] = future - + # Execute await self.handler.reply_received_callback(self.mock_message) - + # Verify self.assertTrue(future.done()) self.assertEqual(future.result(), self.mock_data_response.body()) self.mock_message.ack.assert_called_once_with(multiple=False) - + async def test_reply_received_callback_no_future(self): """Test handling a reply when there's no matching future.""" # Execute await self.handler.reply_received_callback(self.mock_message) - + # Verify self.mock_message.ack.assert_called_once_with(multiple=False) - + def test_ensure_trace_info(self): """Test ensuring trace info is properly set.""" # Execute self.handler.ensure_trace_info(self.mock_data_message) - + # Verify self.tracer.start_span.assert_called_once() - + def test_record_duration(self): """Test recording the duration of message processing.""" # Setup start_time = 1234567880.0 # 10 seconds before mock_time - + # Execute self.handler.record_duration(start_time, 123) - + # No specific assertions needed as this is mostly logging - + def test_collect_trace_info(self): """Test collecting trace info.""" # Execute result = collect_trace_info() - + # Verify it returns a dictionary self.assertIsInstance(result, dict) - + def test_get_default_trace_state(self): """Test getting default trace state.""" # Execute result = get_default_trace_state() - + # Verify it returns a TraceState self.assertIsNotNone(result) - + def test_set_text(self): """Test setting text in a carrier.""" # Setup carrier = {} - + # Execute set_text(carrier, "key", "value") - + # Verify self.assertEqual(carrier["key"], "value") - - + async def test_reconstitute_data_message_stream_none_result(self): """Test reconstituting a message from a stream when the result is None.""" # Setup self.mock_message.headers = {"clevermicro.addressing": 1} self.mock_from_stream.return_value = None - + # Execute result = await self.handler.reconstitute_data_message(self.mock_message) - + # Verify self.assertIsNone(result) self.mock_from_stream.assert_called_once() self.file_handler.on_message.assert_not_called() - + async def test_run_http_request_magic_file_removal_exception(self): """Test handling of an exception when removing files.""" # Setup self.service_adapter.on_message.return_value = self.mock_data_response - multipart_message = MultipartDataMessage(self.mock_data_message, {"file1": "/tmp/file1.txt"}) + multipart_message = MultipartDataMessage( + self.mock_data_message, {"file1": "/tmp/file1.txt"} + ) self.mock_os_remove.side_effect = OSError("File removal error") - + # Execute await self.handler.run_http_request_magic(self.mock_message, multipart_message, self.loop) - + # Verify self.mock_os_remove.assert_called_once_with("/tmp/file1.txt") # Should continue without raising an exception - + async def test_send_reply_content_type_list(self): """Test sending a reply with content type as a list.""" # Setup self.mock_data_response.content_type.return_value = ["application/json", "text/plain"] - + # Execute await self.handler.send_reply("reply-queue", self.mock_data_response) - + # Verify self.mock_run_coroutine.assert_called_once() - # Check that the first content type was used - call_args = self.mock_run_coroutine.call_args[0][0] - # We verify that Message is created with correct content type by inspecting the call args if available. - self.assertEqual(Message.call_args[1]['content_type'], "application/json") + # We verify that Message is created with correct content type by inspecting the RabbitMQ call args if available. + self.assertEqual( + "application/json", self.reply_to_exchange.publish.call_args[1]["message"].content_type + ) if __name__ == "__main__": diff --git a/tests/service/test_amq_service.py b/tests/service/test_amq_service.py index feb2bfa..7f5ff6f 100644 --- a/tests/service/test_amq_service.py +++ b/tests/service/test_amq_service.py @@ -34,55 +34,43 @@ class TestAMQService(unittest.IsolatedAsyncioTestCase): # Check that run_coroutine_threadsafe was called with the correct arguments # We need to patch asyncio.run_coroutine_threadsafe to capture its arguments - with patch('asyncio.run_coroutine_threadsafe') as mock_run: + with patch("asyncio.run_coroutine_threadsafe") as mock_run: self.service.backpressure(current_availability=50, maximum=200) mock_run.assert_called_once() - - # The first argument should be the coroutine - coro_arg = mock_run.call_args[0][0] + # The second argument should be the loop - loop_arg = mock_run.call_args[0][1] - + loop_arg = mock_run.call_args[1]["loop"] + self.assertEqual(loop_arg, self.service.backpressure_handler.loop) - - # For the coroutine, we can check its __qualname__ to verify it's the right method - self.assertEqual( - coro_arg.__qualname__, - 'BackpressureHandler.update_backpressure_value' - ) - + # We can also verify the arguments passed to update_backpressure_value - self.service.backpressure_handler.update_backpressure_value.assert_called_with( - 50, 200 - ) + self.service.backpressure_handler.update_backpressure_value.assert_called_with(50, 200) async def test_backpressure_without_maximum(self): """Test backpressure method without a maximum value.""" # Set a maximum_availability value self.service.maximum_availability = 100 - - with patch('asyncio.run_coroutine_threadsafe') as mock_run: + + with patch("asyncio.run_coroutine_threadsafe") as mock_run: # Call the backpressure method without a maximum value self.service.backpressure(current_availability=50) - + # Verify the handler was called with the right values - self.service.backpressure_handler.update_backpressure_value.assert_called_with( - 50, 100 - ) + self.service.backpressure_handler.update_backpressure_value.assert_called_with(50, 100) + mock_run.assert_called_once() async def test_backpressure_with_negative_maximum(self): """Test backpressure method with a negative maximum value.""" # Set a maximum_availability value self.service.maximum_availability = 100 - - with patch('asyncio.run_coroutine_threadsafe') as mock_run: + + with patch("asyncio.run_coroutine_threadsafe") as mock_run: # Call the backpressure method with a negative maximum value self.service.backpressure(current_availability=50, maximum=-1) - + # Verify the handler was called with the right values - self.service.backpressure_handler.update_backpressure_value.assert_called_with( - 50, 100 - ) + self.service.backpressure_handler.update_backpressure_value.assert_called_with(50, 100) + mock_run.assert_called_once() async def test_backpressure_no_handler(self): """Test backpressure method when no handler is available.""" diff --git a/tests/test_amq_service.py b/tests/test_amq_service.py deleted file mode 100644 index ca97986..0000000 --- a/tests/test_amq_service.py +++ /dev/null @@ -1,74 +0,0 @@ -import asyncio -import unittest -from unittest.mock import AsyncMock, MagicMock, patch - -from amqp.service.amq_service import AMQService - - -class TestAMQService(unittest.IsolatedAsyncioTestCase): - def setUp(self): - # Mock the configuration - self.config = MagicMock() - self.config.amq_adapter.service_name = "test-service" - self.config.amq_adapter.swarm_task_id = "test-task" - self.config.amq_adapter.swarm_service_id = "test-service-id" - self.config.backpressure.threshold_load = 100 - - # Mock the service adapter - self.service_adapter = MagicMock() - - # Patch the get_unique_instance_id function - with patch("amqp.service.amq_service.get_unique_instance_id", return_value=12345): - # Create the AMQService instance - self.service = AMQService(self.config, self.service_adapter) - - # Mock the backpressure handler - self.service.backpressure_handler = MagicMock() - self.service.backpressure_handler.loop = asyncio.get_event_loop() - - async def test_backpressure_with_maximum(self): - """Test backpressure method with a specified maximum value.""" - # Call the backpressure method with a maximum value - self.service.backpressure(current_availability=50, maximum=200) - - # Check that update_backpressure_value was called with the correct values - self.service.backpressure_handler.update_backpressure_value.assert_called_once() - # Get the coroutine that was passed to run_coroutine_threadsafe - coro = self.service.backpressure_handler.update_backpressure_value.call_args[0] - self.assertEqual(coro, (50, 200)) - - async def test_backpressure_without_maximum(self): - """Test backpressure method without a maximum value.""" - # Call the backpressure method without a maximum value - self.service.backpressure(current_availability=50) - - # Check that update_backpressure_value was called with the correct values - self.service.backpressure_handler.update_backpressure_value.assert_called_once() - # Get the coroutine that was passed to run_coroutine_threadsafe - coro = self.service.backpressure_handler.update_backpressure_value.call_args[0] - self.assertEqual(coro, (50, 100)) # Should use threshold_load from config - - async def test_backpressure_with_negative_maximum(self): - """Test backpressure method with a negative maximum value.""" - # Call the backpressure method with a negative maximum value - self.service.backpressure(current_availability=50, maximum=-1) - - # Check that update_backpressure_value was called with the correct values - self.service.backpressure_handler.update_backpressure_value.assert_called_once() - # Get the coroutine that was passed to run_coroutine_threadsafe - coro = self.service.backpressure_handler.update_backpressure_value.call_args[0] - self.assertEqual(coro, (50, 100)) # Should use threshold_load from config - - async def test_backpressure_no_handler(self): - """Test backpressure method when no handler is available.""" - # Set the backpressure handler to None - self.service.backpressure_handler = None - - # Call the backpressure method - self.service.backpressure(current_availability=50) - - # Nothing should happen, no exception should be raised - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_backpressure_handler.py b/tests/test_backpressure_handler.py deleted file mode 100644 index 6fb5e14..0000000 --- a/tests/test_backpressure_handler.py +++ /dev/null @@ -1,253 +0,0 @@ -import asyncio -import time -import unittest -from unittest.mock import AsyncMock, MagicMock, patch - -from amqp.adapter.backpressure_handler import ( - BackpressureHandler, - RunningAverage, - ScaleRequestV1, - ScalingRequestAlert, -) - - -class TestRunningAverage(unittest.TestCase): - def test_running_average_empty(self): - avg = RunningAverage(5) - self.assertEqual(avg.get_average(), 0.0) - - def test_running_average_partial(self): - avg = RunningAverage(5) - avg.add(10) - avg.add(20) - self.assertEqual(avg.get_average(), 15.0) - - def test_running_average_full(self): - avg = RunningAverage(3) - avg.add(10) - avg.add(20) - avg.add(30) - self.assertEqual(avg.get_average(), 20.0) - - def test_running_average_overflow(self): - avg = RunningAverage(3) - avg.add(10) - avg.add(20) - avg.add(30) - avg.add(40) - avg.add(50) - self.assertEqual(avg.get_average(), 40.0) - - def test_get_current_values_partial(self): - avg = RunningAverage(5) - avg.add(10) - avg.add(20) - self.assertEqual(avg.get_current_values(), [10, 20]) - - def test_get_current_values_full(self): - avg = RunningAverage(3) - avg.add(10) - avg.add(20) - avg.add(30) - self.assertEqual(avg.get_current_values(), [10, 20, 30]) - - def test_get_current_values_overflow(self): - avg = RunningAverage(3) - avg.add(10) - avg.add(20) - avg.add(30) - avg.add(40) - avg.add(50) - self.assertEqual(avg.get_current_values(), [30, 40, 50]) - - -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.cpu_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_load = 0 - self.handler.update_current_load(5) - self.assertEqual(self.handler.current_load, 5) - - async def test_update_last_data_message_time(self): - old_time = self.handler.last_backpressure_event_time - self.handler.update_last_data_message_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_load, 50) - self.assertEqual(self.handler.maximum_load, 200) - - # Test updating with a lower maximum (should not change) - await self.handler.update_backpressure_value(60, 50) - self.assertEqual(self.handler.current_load, 60) - self.assertEqual(self.handler.maximum_load, 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.maximum_load = 100 - self.handler.current_load = 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.currentAvailability, 60) - self.assertEqual(args.maxAvailability, 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.maximum_load = 100 - self.handler.current_load = 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.currentAvailability, 10) - self.assertEqual(args.maxAvailability, 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.maximum_load = 100 - self.handler.current_load = 90 - self.handler.last_backpressure_event = ScalingRequestAlert.IDLE - self.handler.last_backpressure_event_time = time.time() - 3 # 3 seconds ago (> idle_duration) - - 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.currentAvailability, 90) - self.assertEqual(args.maxAvailability, 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.maximum_load = -1 - self.handler.current_load = 50 - 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 set maximum to current and send UPDATE - self.assertEqual(self.handler.maximum_load, 50) - mock_publish.assert_called_once() - args = mock_publish.call_args[0][0] - self.assertEqual(args.requestType, ScalingRequestAlert.UPDATE) - self.assertEqual(args.currentAvailability, 50) - self.assertEqual(args.maxAvailability, 50) - - # Test with a higher value that should update the maximum - mock_publish.reset_mock() - await self.handler.update_backpressure_value(80, -1) - - # Should update maximum to 80 - self.assertEqual(self.handler.maximum_load, 80) - mock_publish.assert_called_once() - args = mock_publish.call_args[0][0] - self.assertEqual(args.currentAvailability, 80) - self.assertEqual(args.maxAvailability, 80) - - @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request") - async def test_handle_backpressure_overload_event(self, mock_publish): - self.handler.maximum_load = 100 - self.handler.current_load = 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.currentAvailability, 10) - self.assertEqual(args.maxAvailability, 100) - - @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request") - async def test_handle_backpressure_idle_event(self, mock_publish): - self.handler.maximum_load = 100 - self.handler.current_load = 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.currentAvailability, 90) - self.assertEqual(args.maxAvailability, 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) - - -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() diff --git a/tools/test_amq_service.py b/tools/test_amq_service.py new file mode 100644 index 0000000..31bcac9 --- /dev/null +++ b/tools/test_amq_service.py @@ -0,0 +1,172 @@ +import unittest +from unittest.mock import AsyncMock, Mock, patch + +from amqp.config.amq_configuration import AMQAdapter, AMQConfiguration +from amqp.service.amq_service import AMQService + + +class TestAMQService(unittest.IsolatedAsyncioTestCase): + def setUp(self): + # Create a mock configuration + self.mock_adapter = Mock(spec=AMQAdapter) + self.mock_adapter.service_name = "test-service" + self.mock_adapter.route_mapping = {} + self.mock_adapter.swarm_service_id = "test-swarm-id" + self.mock_adapter.swarm_task_id = "test-task-id" + self.mock_adapter.consul_host = "consul" + self.mock_adapter.consul_port = 4500 + self.mock_adapter.consul_counter_key = "services/ids/counter" + self.mock_adapter.consul_max_retries = 5 + self.mock_adapter.consul_initial_retry_delay = 0.2 + + self.mock_config = Mock(spec=AMQConfiguration) + self.mock_config.amq_adapter = self.mock_adapter + + # Add dispatch configuration + self.mock_dispatch = Mock() + self.mock_dispatch.amq_host = "localhost" + self.mock_dispatch.amq_port = 5672 + self.mock_dispatch.rabbit_mq_user = "guest" + self.mock_dispatch.rabbit_mq_password = "guest" + self.mock_dispatch.use_dlq = False + self.mock_config.dispatch = self.mock_dispatch + + # Add backpressure configuration + self.mock_backpressure = Mock() + self.mock_backpressure.threshold_load = 10 + self.mock_config.backpressure = self.mock_backpressure + + # Create service instance with mocked configuration + self.service = AMQService(self.mock_config) + + # Mock the event loop + self.mock_loop = Mock() + self.service.loop = self.mock_loop + + def tearDown(self): + # Clean up any resources + if hasattr(self, "service"): + self.service.cleanup() + + @patch("amqp.service.amq_service.RabbitMQClient") + @patch("amqp.service.amq_service.DataMessageFactory") + @patch("amqp.service.amq_service.ServiceMessageFactory") + @patch("amqp.service.amq_service.initialize_telemetry") + def test_init( + self, mock_telemetry, mock_service_factory, mock_data_factory, mock_rabbit_client + ): + # Test initialization of AMQService + service = AMQService(self.mock_config) + + # Verify that all required components are initialized + self.assertIsNotNone(service.amq_configuration) + self.assertIsNotNone(service.data_message_factory) + self.assertIsNotNone(service.rabbit_mq_client) + self.assertIsNotNone(service.tracer) + self.assertIsNotNone(service.service_message_factory) + + @patch("amqp.service.amq_service.RabbitMQClient") + async def test_register_routes_valid(self, mock_rabbit_client): + # Setup mock route mapping + self.mock_adapter.route_mapping = "test-service:test-exchange:test-queue:test.#:5:0" + + # Setup mock rabbit client + mock_client = Mock() + mock_client.register_inbound_routes = AsyncMock() + mock_client.register_data_reply_callback = AsyncMock() + mock_client.init = AsyncMock() + mock_client.request_routes = AsyncMock() + mock_client.channel = Mock() + mock_client.channel.is_closed = False + self.service.rabbit_mq_client = mock_client + + # Setup mock service adapter + mock_service_adapter = Mock() + mock_service_adapter.loop = self.mock_loop + + # Setup mock message handler + mock_handler = Mock() + mock_handler.inbound_data_message_callback_no_thread = AsyncMock() + mock_handler.reply_received_callback = AsyncMock() + self.service.amq_data_message_handler = mock_handler + + # Initialize the service + await self.service.init(self.mock_loop, mock_service_adapter) + + # Call register_routes + await self.service.register_routes() + + # Verify that routes were registered twice (once in init, once in register_routes) + self.assertEqual(mock_client.register_inbound_routes.call_count, 2) + self.assertEqual(mock_client.register_data_reply_callback.call_count, 2) + + @patch("amqp.service.amq_service.RabbitMQClient") + async def test_register_routes_invalid(self, mock_rabbit_client): + # Setup invalid route mapping + self.mock_adapter.route_mapping = None + + # Setup mock rabbit client + mock_client = Mock() + mock_client.register_inbound_routes = AsyncMock() + mock_client.register_data_reply_callback = AsyncMock() + self.service.rabbit_mq_client = mock_client + + # Call register_routes + result = await self.service.register_routes() + + # Verify that no routes were registered + mock_client.register_inbound_routes.assert_not_called() + mock_client.register_data_reply_callback.assert_not_called() + self.assertIsNone(result) + + @patch("amqp.service.amq_service.RabbitMQClient") + async def test_reinitialize_if_disconnected(self, mock_rabbit_client): + # Setup mock route mapping + self.mock_adapter.route_mapping = "test-service:test-exchange:test-queue:test.#:5:0" + + # Setup mock rabbit client with closed channel + mock_client = AsyncMock() + mock_client.channel = Mock() + mock_client.channel.is_closed = True + mock_client.init = AsyncMock() + mock_client.register_inbound_routes = AsyncMock() + mock_client.register_data_reply_callback = AsyncMock() + mock_client.request_routes = AsyncMock() + + # Make the constructor return our mock client + mock_rabbit_client.return_value = mock_client + self.service.rabbit_mq_client = mock_client + + # Setup mock message handler + mock_handler = Mock() + mock_handler.inbound_data_message_callback_no_thread = AsyncMock() + mock_handler.reply_received_callback = AsyncMock() + self.service.amq_data_message_handler = mock_handler + + # Call reinitialize_if_disconnected + await self.service.reinitialize_if_disconnected() + + # Verify that client was reinitialized + mock_client.init.assert_called_once() + + def test_remove_outstanding_future(self): + # Setup mock message handler with outstanding message + self.service.amq_data_message_handler = Mock() + self.service.amq_data_message_handler.outstanding = {"test-id": Mock()} + + # Remove future + self.service.remove_outstanding_future("test-id") + + # Verify message was removed from outstanding + self.assertNotIn("test-id", self.service.amq_data_message_handler.outstanding) + + def test_cleanup(self): + # Setup mock rabbit client + mock_client = Mock() + self.service.rabbit_mq_client = mock_client + + # Call cleanup + self.service.cleanup() + + # Verify rabbit client cleanup was called + mock_client.cleanup.assert_called_once() diff --git a/tools/test_backpressure_handler.py b/tools/test_backpressure_handler.py new file mode 100644 index 0000000..19e9feb --- /dev/null +++ b/tools/test_backpressure_handler.py @@ -0,0 +1,736 @@ +# +# 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()