feat/58 - simplify complex logic
This commit is contained in:
@@ -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,54 +90,42 @@ 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
|
||||
)
|
||||
if maximum > 0 and maximum > self.maximum_load:
|
||||
self.maximum_load = maximum
|
||||
|
||||
# Update the current parallel executions count
|
||||
self.update_current_load(current)
|
||||
|
||||
# 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,
|
||||
"Backpressure: Updating backpressure value, current_availability=%s, maximum=%s",
|
||||
current_availability,
|
||||
maximum,
|
||||
)
|
||||
self.maximum_load = maximum
|
||||
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()
|
||||
|
||||
def start_backpressure_monitor(self) -> Optional[Thread]:
|
||||
# Start the Backpressure monitor loop
|
||||
@@ -192,7 +136,7 @@ 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.
|
||||
@@ -200,26 +144,22 @@ class BackpressureHandler:
|
||||
- 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
|
||||
@@ -227,7 +167,7 @@ class BackpressureHandler:
|
||||
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 (
|
||||
@@ -245,16 +185,26 @@ class BackpressureHandler:
|
||||
# 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
|
||||
@@ -262,95 +212,13 @@ class BackpressureHandler:
|
||||
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 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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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__))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@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
|
||||
|
||||
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.current_availability, 10)
|
||||
self.assertEqual(args.max_availability, 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.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
|
||||
)
|
||||
|
||||
# Call the method
|
||||
await handler.check_overload_condition()
|
||||
await self.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
|
||||
# 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)
|
||||
|
||||
@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_load = 2
|
||||
handler.maximum_load = 10
|
||||
handler.last_backpressure_event = ScalingRequestAlert.IDLE # Coming from IDLE state
|
||||
|
||||
with patch("time.time") as mock_time:
|
||||
# Set current 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
|
||||
@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
|
||||
)
|
||||
|
||||
# Call the method
|
||||
await handler.check_overload_condition()
|
||||
await self.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
|
||||
# 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)
|
||||
|
||||
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
|
||||
# Test with a higher value that should update the maximum
|
||||
mock_publish.reset_mock()
|
||||
await self.handler.update_backpressure_value(80, 80)
|
||||
|
||||
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
|
||||
thread = handler.start_backpressure_monitor()
|
||||
# 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)
|
||||
|
||||
# Verify no thread was started
|
||||
assert thread is not None
|
||||
assert hasattr(handler, "thread") or handler.thread is not None
|
||||
@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
|
||||
|
||||
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
|
||||
await self.handler.handle_backpressure_overload_event()
|
||||
|
||||
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()
|
||||
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)
|
||||
|
||||
# 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()
|
||||
@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()
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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):
|
||||
@@ -66,7 +70,7 @@ 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
|
||||
@@ -84,10 +88,12 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
|
||||
# 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
|
||||
@@ -96,35 +102,43 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
|
||||
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):
|
||||
@@ -163,7 +177,7 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
# 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."""
|
||||
@@ -179,7 +193,7 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
|
||||
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
|
||||
|
||||
@@ -208,7 +222,7 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
|
||||
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()
|
||||
|
||||
@@ -218,7 +232,9 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
|
||||
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()
|
||||
@@ -231,7 +247,9 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
|
||||
"""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)
|
||||
@@ -245,7 +263,9 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
|
||||
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()
|
||||
@@ -258,10 +278,12 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
|
||||
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):
|
||||
@@ -347,7 +369,6 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
|
||||
# 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
|
||||
@@ -366,7 +387,9 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
|
||||
"""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
|
||||
@@ -386,10 +409,10 @@ class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
# 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__":
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user