clevermicro/amq-adapter-python
14
0
Fork 0
Code Issues 16 Pull Requests 4 Actions Packages Projects Releases Wiki Activity

Implement unit test #47

Merged
stanislav.hejny merged 11 commits from unittest-coverage-85 into develop 2025-05-19 14:30:16 +00:00
Conversation 2 Commits 11 Files Changed 29
+3269 -421
29 changed files with 3269 additions and 421 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.cursor/rules/python.mdc
+118
View File
@@ -0,0 +1,118 @@
---
description:
globs:
alwaysApply: true
---
---
description: Python best practices and patterns for modern software development with Flask and SQLite
globs: **/*.py, src/**/*.py, tests/**/*.py
---
# Python Best Practices
## Project Structure
- Use src-layout with `amqp/`
- Place tests in `tests/` directory parallel to `amqp/`
- Keep configuration in `amqp/config/` or as environment variables
- Store requirements in `pyproject.toml`
- Place static files in `static/` directory
- Use `templates/` for Jinja2 templates
## Code Style
- Follow Black code formatting
- Use isort for import sorting
- Follow PEP 8 naming conventions:
- snake_case for functions and variables
- PascalCase for classes
- UPPER_CASE for constants
- Maximum line length of 100 characters
- Use absolute imports over relative imports
## Type Hints
- Use type hints for all function parameters and returns
- Import types from `typing` module
- Use `Optional[Type]` instead of `Type | None`
- Use `TypeVar` for generic types
- Define custom types in `types.py`
- Use `Protocol` for duck typing
## Database
- Use SQLAlchemy ORM
- Implement database migrations with Alembic
- Use proper connection pooling
- Define models in separate modules
- Implement proper relationships
- Use proper indexing strategies
## Authentication
##- Use Flask-Login for session management
##- Implement Google OAuth using Flask-OAuth
##- Hash passwords with bcrypt
##- Use proper session security
##- Implement CSRF protection
##- Use proper role-based access control
## API Design
##- Use Flask-RESTful for REST APIs
##- Implement proper request validation
##- Use proper HTTP status codes
##- Handle errors consistently
##- Use proper response formats
##- Implement proper rate limiting
## Testing
- Use pytest for testing
- Write tests for all routes
- Use pytest-cov for coverage
- Implement proper fixtures
- Use proper mocking with pytest-mock
- Test all error scenarios
## Security
##- Use HTTPS in production
##- Implement proper CORS
##- Sanitize all user inputs
##- Use proper session configuration
- Implement proper logging
##- Follow OWASP guidelines
## Performance
##- Use proper caching with Flask-Caching
##- Implement database query optimization
##- Use proper connection pooling
##- Implement proper pagination
##- Use background tasks for heavy operations
- Monitor application performance
## Error Handling
- Create custom exception classes
- Use proper try-except blocks
- Implement proper logging
- Return proper error responses
- Handle edge cases properly
- Use proper error messages
## Documentation
- Use Google-style docstrings
- Document all public APIs
- Keep README.md updated
- Use proper inline comments
- Generate API documentation
- Document environment setup
## Development Workflow
- Use virtual environments (venv)
- Use or Implement pre-commit hooks
- Use proper Git workflow
- Follow semantic versioning
- Use proper CI/CD practices
- Implement proper logging
## Dependencies
##- Pin dependency versions
##- Use requirements.txt for production
- Separate dev dependencies
##- Use proper package versions
##- Regularly update dependencies
- Check for security vulnerabilitie
]cvx -=
amqp/adapter/backpressure_handler.py
+107 -81
View File
@@ -48,8 +48,15 @@ class ScaleRequestV1:
class RunningAverage:
def __init__(self, num_items):
self.buffer = [0] * num_items # Fixed-size array
"""
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
@@ -89,6 +96,8 @@ class BackpressureHandler:
# helps prevent flooding the system with backpressure events
last_backpressure_event_time = 0
last_backpressure_event = ScalingRequestAlert.UPDATE
# _callback_list is used in unit tests to record the invoked callbacks
_callback_list = None
def __init__(
self,
@@ -105,6 +114,18 @@ class BackpressureHandler:
self.swarm_task_id = self.config.amq_adapter.swarm_task_id
self.parallel_workers = self.config.backpressure.threshold_threads
self.average_cpu_usage = None
self.do_loop = -1
def _do_loop(self) -> bool:
"""
Check if the loop should continue running.
Positive value of do_loop indicates the number of iterations left.
Negative value indicates the loop should run indefinitely.
"""
_val = self.do_loop != 0
if self.do_loop > 0:
self.do_loop -= 1
return _val
def increase_parallel_executions(self):
"""Increase the number of parallel executions"""
@@ -135,6 +156,7 @@ class BackpressureHandler:
def start_backpressure_monitor(self) -> Thread:
# Start the Backpressure monitor loop
self.thread = Thread(target=self.backpressure_monitor_loop)
self.thread.daemon = True # This makes it a daemon thread
self.thread.start()
return self.thread
@@ -154,93 +176,93 @@ class BackpressureHandler:
self.last_backpressure_event_time = time.time()
self.last_backpressure_event = ScalingRequestAlert.OVERLOAD
def backpressure_monitor_loop(self):
async def _backpressure_monitor(self):
"""Monitor the backpressure conditions and trigger events accordingly"""
_time_window_millis = self.config.backpressure.time_window
_idle_duration_millis = self.config.backpressure.idle_duration
_overload_duration_millis = self.config.backpressure.idle_duration
_loop = asyncio.new_event_loop()
_monitor_interval = 0.2 # Monitor every 200ms
_overload_duration_millis = self.config.backpressure.cpu_overload_duration
_monitor_interval = 0.5 # Monitor every 500ms
CPU_IDLE, CPU_ACTIVE, CPU_OVERLOAD = 0, 1, 2
IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 15, 90
self.average_cpu_usage = RunningAverage(
int(max(_overload_duration_millis, _idle_duration_millis) // _monitor_interval)
)
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():
_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_data_message_time,
# 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
async def _backpressure_monitor():
"""Monitor the backpressure conditions and trigger events accordingly"""
_cpu_usage_state = CPU_ACTIVE
_cpu_usage_changed = time.time()
while True:
_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_data_message_time,
# 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
# 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_data_message_time > _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,
)
):
# Trigger the overload event
await self.handle_backpressure_overload_event()
# 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.OVERLOAD
elif (
_cpu_usage_state == CPU_IDLE
and _now - _cpu_usage_changed > _idle_duration_millis
and _now - self.last_data_message_time > _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
await asyncio.sleep(_monitor_interval)
self.last_backpressure_event = ScalingRequestAlert.UPDATE
await asyncio.sleep(_monitor_interval)
_loop.run_until_complete(_backpressure_monitor())
def backpressure_monitor_loop(self):
_loop = asyncio.new_event_loop()
_loop.run_until_complete(self._backpressure_monitor())
async def handle_backpressure_overload_event(self):
logging_warning("Backpressure: Capacity close to depleted!")
@@ -283,6 +305,8 @@ class BackpressureHandler:
_exchange = await channel.get_exchange(name="cleverthis.clevermicro.management")
return _exchange
if BackpressureHandler._callback_list is not None:
BackpressureHandler._callback_list["_wrap_rabbit_mq_api_init"] = 1
self.exchange = await await_result(
asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api_init(self.channel), self.loop)
)
@@ -312,4 +336,6 @@ class BackpressureHandler:
return True
return False
if BackpressureHandler._callback_list is not None:
BackpressureHandler._callback_list["_wrap_rabbit_mq_api"] = 1
await await_future(asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api(), self.loop))
amqp/adapter/cleverthis_service_adapter.py
+25 -22
View File
@@ -36,6 +36,30 @@ from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE
from amqp.router.utils import await_result
async def _create_reply_exchange_and_queue(
connection: AbstractRobustConnection,
) -> tuple[AbstractRobustChannel, AbstractExchange, str]:
"""
Create objects for reply.
Return (channel for reply queue, exchange for reply, reply queue name)
"""
# Wrap calls to RabbitMQ API to make it Awaitable, so it can be run thread-safe with
# correct event loop. This is needed because the RabbitMQ API is not thread-safe and needs
# to be called in the context of event loop the first RabbitMQ connection was created with.
_channel = await connection.channel()
_exchange_name: str = AMQ_REPLY_TO_EXCHANGE
_exchange = await _channel.get_exchange(name=_exchange_name, ensure=True)
_queue = await _channel.declare_queue(exclusive=False) # Declare a unique, non-exclusive queue
_queue_name = _queue.name
await _queue.bind(
exchange=_exchange_name,
routing_key=_queue_name, # Use queue name as routing key
)
logging_debug(f"Publishing to exchange: {_exchange_name}, q: {_queue_name}")
return _channel, _exchange, _queue_name
async def _convert_file_response(
obj: FileResponse,
connection: AbstractRobustConnection,
@@ -47,33 +71,12 @@ async def _convert_file_response(
The file content itself is transmitted through the named stream (which is implemented as
temporary dedicated queue.
"""
async def _wrap_amq_api_calls(
connection: AbstractRobustConnection,
) -> tuple[AbstractRobustChannel, AbstractExchange, str]:
# Wrap calls to RabbitMQ API to make it Awaitable, so it can be run thread-safe with
# correct event loop. This is needed because the RabbitMQ API is not thread-safe and needs
# to be called in the context of event loop the first RabbitMQ connection was created with.
_channel = await connection.channel()
_exchange_name: str = AMQ_REPLY_TO_EXCHANGE
_exchange = await _channel.get_exchange(name=_exchange_name, ensure=True)
_queue = await _channel.declare_queue(
exclusive=False
) # Declare a unique, non-exclusive queue
_queue_name = _queue.name
await _queue.bind(
exchange=_exchange_name,
routing_key=_queue_name, # Use queue name as routing key
)
logging_debug(f"Publishing to exchange: {_exchange_name}, q: {_queue_name}")
return _channel, _exchange, _queue_name
# 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
# 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(
asyncio.run_coroutine_threadsafe(_wrap_amq_api_calls(connection), loop)
asyncio.run_coroutine_threadsafe(_create_reply_exchange_and_queue(connection), loop)
)
# publish_file also needs RabbitMQ API, ensure the loop is passed as argument
(_stream_name, _file_size) = await file_handler.publish_file(
amqp/adapter/data_parser.py
+3 -4
View File
@@ -40,7 +40,7 @@ class MultipartFormDataParser:
try:
self.form_data = json.loads(message.body())
self.is_json = True
if self.form_data["files"] and self.form_data["form"]:
if self.form_data["files"] is not None and self.form_data["form"] is not None:
self.form_data = AMQDataParser.process_form_data(self.form_data)
except Exception as jsonerror:
logging_error(f"Parsing JSON: {jsonerror}")
@@ -55,7 +55,7 @@ class MultipartFormDataParser:
try:
self.form_data = json.loads(message.body())
self.is_json = True
if self.form_data["files"] and self.form_data["form"]:
if self.form_data["files"] is not None and self.form_data["form"] is not None:
self.form_data = AMQDataParser.process_form_data(self.form_data)
except Exception as jsonerror:
logging_error(f"Parsing JSON: {jsonerror}")
@@ -109,7 +109,6 @@ class AMQDataParser:
for key, value in data["files"].items():
if value: # Only override if files data is non-empty
result[key] = value[0]
return result
@staticmethod
@@ -162,7 +161,7 @@ class AMQDataParser:
}
if len(_form_data) == 0 and len(message.body()) > 0:
_form_data = json.loads(message.body())
if _form_data["files"] and _form_data["form"]:
if _form_data["files"] is not None and _form_data["form"] is not None:
_form_data = AMQDataParser.process_form_data(_form_data)
return _form_data
except Exception as e:
amqp/adapter/file_handler.py
+74 -50
View File
@@ -153,6 +153,44 @@ class StreamingFileHandler(FileHandler):
def __init__(self):
super().__init__()
# This method MUST NOT be static.
# For testing purpose, we will mock this method and inject a mock connection.
async def _get_connection(self, rabbitmq_url, loop):
return await connect_robust(rabbitmq_url, loop=loop)
async def _download_from_queue_to_file(
self, rabbitmq_url, queue_name, filepath, loop
) -> tuple[int, int] | None:
"""
Internal method for downloading a file from a queue to disk.
Return (file size, eof status).
"""
# Here the normal mock doesn't work since this method
# is async and will be run in another loop,
# out of the scope of method mock.
connection = await self._get_connection(rabbitmq_url, loop=loop)
_file_size = 0
_eof = IN_PROGRESS
try:
channel: AbstractRobustChannel = await connection.channel()
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
logging_debug(f"Awaiting all file chunks being downloaded for {filepath}")
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
with open(filepath, "ab") as f:
f.write(message.body)
_file_size = _file_size + message.body_size
_eof = int(message.headers.get("eof", END_OF_FILE))
if _eof != IN_PROGRESS:
break
except Exception as e:
logging_error(f"File download failed: {e}")
finally:
await connection.close()
return _file_size, _eof
async def download_file(
self, loop, rabbitmq_url, file_info, message_data, output_dir
) -> (int, int):
@@ -184,39 +222,46 @@ class StreamingFileHandler(FileHandler):
f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name}:{queue_name} to {filepath}"
)
async def wrap_rabbit_mq_calls() -> tuple[int, int] | None:
connection = await connect_robust(rabbitmq_url, loop=loop)
try:
_file_size = 0
_eof = IN_PROGRESS
channel: AbstractRobustChannel = await connection.channel()
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
logging_debug(f"Awaiting all file chunks being downloaded for {filename}")
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
with open(filepath, "ab") as f:
f.write(message.body)
_file_size = _file_size + message.body_size
_eof = int(message.headers.get("eof", END_OF_FILE))
if _eof != IN_PROGRESS:
break
except Exception as e:
logging_error(f"File download failed: {e}")
finally:
await connection.close()
return _file_size, _eof
if queue_name:
_local_file_size, _local_eof = await await_result(
asyncio.run_coroutine_threadsafe(wrap_rabbit_mq_calls(), loop=loop)
asyncio.run_coroutine_threadsafe(
self._download_from_queue_to_file(rabbitmq_url, queue_name, filepath, loop),
loop=loop,
)
)
return _local_file_size, _local_eof
return 0, CANCELLED
async def download_buffer(self, loop, rabbitmq_url, queue_name) -> (bytearray, int):
async def _download_from_queue_to_buffer(
self, rabbitmq_url, queue_name, loop
) -> tuple[bytearray, int]:
"""
Internal method for downloading a file from a queue to RAM.
Return (bytearray, eof status).
"""
_connection = await self._get_connection(rabbitmq_url, loop=loop)
_eof = IN_PROGRESS
_res = bytearray()
try:
channel: AbstractRobustChannel = await _connection.channel()
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
logging_debug(f"Awaiting all buffer chunks being downloaded for {queue_name}")
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
_res.extend(message.body)
_eof = message.headers.get("eof", END_OF_FILE)
if _eof != IN_PROGRESS:
break
logging_info(f"Finished downloading buffer {queue_name}")
except Exception as e:
logging_error(f"Downloading buffer: {e}")
await _connection.close()
return _res, _eof
async def download_buffer(self, loop, rabbitmq_url, queue_name: str | None) -> (bytearray, int):
"""
Asynchronously downloads a bytearay from RabbitMQ based on the provided file information.
It assumes entire array is in the first (and only) chunk
@@ -227,31 +272,10 @@ class StreamingFileHandler(FileHandler):
"""
logging_debug(f"Downloading buffer from stream: {queue_name}")
if queue_name:
async def wrap_rabbit_mq_calls() -> tuple[bytearray, int]:
_connection = await connect_robust(rabbitmq_url, loop=loop)
_eof = IN_PROGRESS
_res = bytearray()
try:
channel: AbstractRobustChannel = await _connection.channel()
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
logging_debug(f"Awaiting all buffer chunks being downloaded for {queue_name}")
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
_res.extend(message.body)
_eof = message.headers.get("eof", END_OF_FILE)
if _eof != IN_PROGRESS:
break
logging_info(f"Finished downloading buffer {queue_name}")
except Exception as e:
logging_error(f"Downloading buffer: {e}")
await _connection.close()
return _res, _eof
_local_res, _local_eof = await await_result(
asyncio.run_coroutine_threadsafe(wrap_rabbit_mq_calls(), loop=loop)
asyncio.run_coroutine_threadsafe(
self._download_from_queue_to_buffer(rabbitmq_url, queue_name, loop), loop=loop
)
)
return _local_res, _local_eof
return 0, CANCELLED
amqp/adapter/logging_utils.py
+1 -1
View File
@@ -120,7 +120,7 @@ def logging_debug(msg: str, *args) -> None:
_thread_id = threading.get_ident()
module, line_number, function_name = get_context_info()
if module and line_number and function_name:
logging.info(
logging.debug(
f"{_thread_id}:{__version__} [DBG] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
*args,
)
amqp/adapter/pydantic_serializer.py
+60 -80
View File
@@ -2,10 +2,12 @@
This module provides serialization and deserialization functions for Pydantic models that extend BaseModel
"""
import importlib
import json
import uuid
from datetime import datetime
from typing import Any, Dict
from types import UnionType
from typing import Any, Dict, Union, get_args, get_origin
from urllib.parse import parse_qs, urlparse
from pydantic import BaseModel
@@ -26,8 +28,12 @@ def serialize_object(obj: BaseModel) -> str:
A string representation of the object.
"""
class_name = obj.__class__.__name__
module_name = obj.__class__.__module__
obj_dict = {}
for key, value in obj.model_dump().items():
for key, _ in obj.model_dump().items():
# here we use the original value
# the one returned by model_dump will convert everything into dict recursively
value = getattr(obj, key)
if isinstance(value, uuid.UUID):
obj_dict[key] = str(value)
elif isinstance(value, datetime):
@@ -47,8 +53,21 @@ def serialize_object(obj: BaseModel) -> str:
else:
obj_dict[key] = value
dict_string = ",".join(f"{k}:{v}" for k, v in obj_dict.items())
return f"{class_name}:{{{dict_string}}}"
dict_string = json.dumps(obj_dict)
return f"{module_name}.{class_name}:{dict_string}"
def try_deserialize_object(data: str) -> BaseModel | str:
"""
Try deserialize object, if failed, return the input string as-is.
"""
if not isinstance(data, str):
return data
try:
return deserialize_object(data)
except Exception as e:
logging_debug(f"Failed to deserialize object: {e}")
return data
def deserialize_object(data: str) -> BaseModel:
@@ -63,93 +82,54 @@ def deserialize_object(data: str) -> BaseModel:
Returns:
An R2RSerializable object.
"""
class_name, dict_str = data.split(":", 1)
dict_str = dict_str.strip("{}")
pairs = dict_str.split(",") if dict_str else []
class_path, dict_str = data.split(":", 1)
module_path, class_name = class_path.rsplit(".", 1)
obj_dict: Dict[str, Any] = {}
for pair in pairs:
if not pair or ":" not in pair:
continue
k, v = pair.split(":", 1)
v = v.strip() # important
obj_dict[k] = _convert_value(v)
obj_dict: Dict[str, Any] = json.loads(dict_str)
logging_debug(f"Deserialized object: {class_name}")
logging_debug(f"Deserialized object: {class_path}")
if "[" in class_name:
class_name = class_name.split("[")[0]
logging_debug(f"trying object: {class_name}")
cls = globals().get(class_name)
module = importlib.import_module(module_path)
cls = getattr(module, class_name)
if cls is None:
raise ValueError(f"Unknown class: {class_name}")
# process nested model, dict and list
for key, value in obj_dict.items():
key_type: type[Any] = cls.model_fields[key].annotation
# handle union with none, this will give the true type
# for example, return `A` if the type is `A | None`
if get_origin(key_type) is UnionType or get_origin(key_type) is Union:
union_types = get_args(key_type)
if len(union_types) > 2:
# if there are multiple unions, like `A | B | None`
# then we have no way to figure out which one to use
raise ValueError(f"Unsupported Union type: {key_type}")
# return the none-None one as the type for processing
if union_types[0] is type(None):
key_type = union_types[1]
else:
key_type = union_types[0]
# skip if the value is None
if value is None:
continue
if key_type == uuid.UUID:
obj_dict[key] = uuid.UUID(value)
elif key_type == datetime:
obj_dict[key] = datetime.fromisoformat(value)
elif key_type == dict:
obj_dict[key] = value
elif issubclass(key_type, BaseModel):
obj_dict[key] = deserialize_object(value)
elif key_type == list:
obj_dict[key] = list(map(try_deserialize_object, value))
return cls(**obj_dict)
def _convert_value(v: str) -> Any:
"""
Helper function to convert a string to its appropriate Python type.
"""
if v == "None":
return None
elif _is_uuid_string(v):
return uuid.UUID(v)
elif _is_datetime_string(v):
return datetime.fromisoformat(v)
elif _is_int_string(v):
return int(v)
elif (
v.startswith("{")
and v.endswith("}")
and ":" in v
and "{" not in v[1:-1]
and "}" not in v[1:-1]
): # nested object
# check if it is a json dict
try:
json.loads(v)
return json.loads(v)
except json.JSONDecodeError:
return deserialize_object(v) # recursive call
elif v.startswith("[") and v.endswith("]"): # list
v = v[1:-1] # remove brackets
items = v.split(",") if v else []
deserialized_list = []
for item in items:
item = item.strip()
deserialized_list.append(_convert_value(item)) # recursive call
return deserialized_list
else:
return v # Keep as string
def _is_uuid_string(s: str) -> bool:
"""Helper to check if a string is a valid UUID."""
try:
uuid.UUID(s)
return True
except ValueError:
return False
def _is_datetime_string(s: str) -> bool:
"""Helper to check if a string is a valid datetime in ISO format"""
try:
datetime.fromisoformat(s)
return True
except ValueError:
return False
def _is_int_string(s: str) -> bool:
"""Helper to check if a string is a valid int."""
try:
int(s)
return True
except ValueError:
return False
def parse_url_query(url):
"""
Parses a HTTP GET URL and returns (context_path, query_params_dict)
amqp/adapter/serializer.py
+5 -2
View File
@@ -18,7 +18,7 @@ def bytes_to_long(b: bytes) -> int:
def read_long(bis: BytesIO) -> int:
try:
return struct.unpack("q", bis.read(8))[0]
except IOError:
except struct.error:
return 0
@@ -51,7 +51,10 @@ def csv_as_list(input_str: str) -> List[str]:
def add_to_map(map_: Dict[str, List[str]], token: str) -> None:
eq_idx = token.index("=")
try:
eq_idx = token.index("=")
except ValueError:
eq_idx = -1 # set to -1 when not found
if eq_idx > 0:
key = token[:eq_idx]
map_[key] = csv_as_list(token[eq_idx + 1 :])
amqp/healthcheck.py
-12
View File
@@ -1,12 +0,0 @@
import sys
import requests
url = "http://localhost:8080/amq-adapter-healthcheck"
headers = {"Authorization": "Bearer my_access_token"}
params = {"limit": 10, "offset": 20}
response = requests.get(url, headers=headers, params=params)
if response.status_code == 200:
sys.exit(0)
# Failed execution
sys.exit(1)
amqp/service/amq_service.py
+7 -2
View File
@@ -17,7 +17,7 @@ from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import DataMessage
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient
from amqp.service.amq_message_handler import DataMessageHandler
from amqp.service.tracing import initialize_trace
from amqp.service.tracing import TracingConfig, initialize_telemetry
logging.basicConfig(level=logging.DEBUG)
logging.info(f"CWD = {os.getcwd()}")
@@ -46,7 +46,12 @@ class AMQService:
self.amq_configuration.amq_adapter.generator_id
)
self.rabbit_mq_client = RabbitMQClient(self.amq_configuration)
self.tracer = initialize_trace(name=amq_configuration.amq_adapter.service_name)
self.tracer = initialize_telemetry(
tracing_config=TracingConfig(
service_name=amq_configuration.amq_adapter.service_name,
service_version=os.getenv("SERVICE_VERSION", "unknown"),
)
)
self.amq_data_message_handler = None
self.service_message_factory = ServiceMessageFactory(
self.amq_configuration.amq_adapter.service_name
amqp/service/tracing.py
+182 -85
View File
@@ -1,99 +1,196 @@
"""
OpenTelemetry configuration for tracing and metrics.
This module provides functionality to initialize and configure OpenTelemetry
for both tracing (Jaeger via OTLP) and metrics (Prometheus) in the AMQ adapter.
"""
import logging
import os
from dataclasses import dataclass
from typing import Any, Dict, Optional
from opentelemetry import metrics, trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.exporter.prometheus import PrometheusMetricReader
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
from opentelemetry.sdk.trace import Tracer, TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from opentelemetry.sdk.resources import SERVICE_NAME, SERVICE_VERSION, Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
from opentelemetry.trace import Tracer
from prometheus_client import start_http_server
logger = logging.getLogger(__name__)
def initialize_trace(app=None, name=None) -> Tracer:
# Resource can be required for some backends, e.g. Jaeger or Prometheus
# If resource wouldn't be set - traces wouldn't appear in Jaeger.
# The 'name' value is the name shown in Jaeger.
resource = Resource(
attributes={
SERVICE_NAME: (name if name else os.environ.get("APPLICATION_NAME", "Python App"))
}
)
#
# 1. Tracing (Jaeger) configuration
# ----------------------------------
#
# Configure the OTLP exporter (or HTTP exporter the same way)
# the recommended configuration is via environment variables. The mandatory ones:
# OTEL_EXPORTER_OTLP_ENDPOINT: "http://jaeger:4317"
# OTEL_EXPORTER_OTLP_INSECURE: "true"
# for more see here: https://opentelemetry-python.readthedocs.io/en/latest/exporter/otlp/otlp.html
# alternatively these values can be provided as input values to OTLPSpanExporter() constructor.
#
# otlp_exporter = OTLPSpanExporter()
# or provide endpoint explicitly, as
otlp_exporter = OTLPSpanExporter(endpoint="http://jaeger:4317")
#
# Initialize the tracer provider. This is abstract API call.
# The actual implementation is given by the included dependency
# The opentelemetry default is the OTLP trace provider.
@dataclass
class TracingConfig:
"""Configuration for OpenTelemetry tracing."""
service_name: str
service_version: str
otlp_endpoint: str = "http://localhost:4317" # Default Jaeger OTLP endpoint
insecure: bool = True
debug: bool = False
additional_attributes: Optional[Dict[str, Any]] = None
@dataclass
class MetricsConfig:
"""Configuration for OpenTelemetry metrics."""
prometheus_port: int = 9464
prometheus_host: str = "localhost"
export_interval_millis: int = 30000
additional_attributes: Optional[Dict[str, Any]] = None
def create_resource(config: TracingConfig) -> Resource:
"""Create an OpenTelemetry resource with service information."""
attributes = {
SERVICE_NAME: config.service_name,
SERVICE_VERSION: config.service_version,
}
if config.additional_attributes:
attributes.update(config.additional_attributes)
return Resource.create(attributes)
def setup_otlp_exporter(config: TracingConfig) -> OTLPSpanExporter:
"""Configure and create an OTLP exporter for Jaeger."""
return OTLPSpanExporter(endpoint=config.otlp_endpoint, insecure=config.insecure)
def setup_trace_provider(config: TracingConfig, resource: Resource) -> TracerProvider:
"""Set up the TracerProvider with OTLP exporter."""
provider = TracerProvider(resource=resource)
# Add the exporters to the batch span processor
span_processor = BatchSpanProcessor(otlp_exporter)
provider.add_span_processor(span_processor)
# Sets the global default tracer provider
trace.set_tracer_provider(provider)
#
# logging.info(' * Instrumenting application..... ')
# if Flask, then use:
# FlaskInstrumentor().instrument_app(flask_app)
# if FastAPI, then use
# FastAPIInstrumentor.instrument_app(app)
# Add OTLP exporter for Jaeger
otlp_exporter = setup_otlp_exporter(config)
provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
#
# Python Flask is now automatically instrumented to export details of each REST endpoint call.
# For Django based app, use 'opentelemetry-instrumentation-django' import instead. see here:
# https://opentelemetry-python.readthedocs.io/en/latest/examples/django/README.html for example
#
# If the used library has no instrumentation support in opentelemetry.io, then the trace
# can be started manually. In this code,
# 'trace' should be global shared variable, initialized once per application (see tracing.py)
# 'tracer' can be shared global or local variable, application can use multiple uniquely
# named tracers, Jaeger will group the generated traces by the name of the tracer
# that was used to start the Span. See this example code:
#
# tracer = trace.get_tracer(name)
#
# @app.route("/server_request")
# def server_request():
# with tracer.start_as_current_span(
# "server_request",
# context=extract(request.headers),
# kind=trace.SpanKind.SERVER,
# attributes=collect_request_attributes(request.environ),
# ):
# # do the work here, like logging.info something and return response
# logging.info(request.args.get("param"))
# return "served"
#
# Note:
# the otlp code looks for HTTP header named 'traceparent' that, if present, should contain
# the otlp parent span identification, and the span created here will be made a child
# of that parent span, indicating visually the progression of requests via different services
# for as long as the current span is provided in the traceparent header to next service.
# Add console exporter in debug mode
if config.debug:
provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))
# 2. Metrics (Prometheus) configuration
# -------------------------------------
#
# Start Prometheus client
start_http_server(port=9464, addr="localhost")
# Initialize PrometheusMetricReader which pulls metrics from the SDK
# on-demand to respond to scrape requests (unlike Jaeger, which worls with push model - application
# pushes traces when those become available to Jaeger's collector. Prometheus works with pull model
# where it periodically polls pre-configured and/or discovered endpoints for metrics.
reader = PrometheusMetricReader()
provider = MeterProvider(resource=resource, metric_readers=[reader])
metrics.set_meter_provider(provider)
return provider
return trace.get_tracer(name)
def setup_metrics_provider(config: MetricsConfig, resource: Resource) -> MeterProvider:
"""Set up the MeterProvider with Prometheus export."""
try:
# Start Prometheus HTTP server
start_http_server(port=config.prometheus_port, addr=config.prometheus_host)
logger.info(
f"Started Prometheus metrics server on {config.prometheus_host}:{config.prometheus_port}"
)
# Create Prometheus reader
prometheus_reader = PrometheusMetricReader()
# Create periodic reader for custom metrics
periodic_reader = PeriodicExportingMetricReader(
prometheus_reader, export_interval_millis=config.export_interval_millis
)
# Create and configure meter provider
return MeterProvider(resource=resource, metric_readers=[prometheus_reader, periodic_reader])
except Exception as e:
logger.error(f"Failed to set up metrics provider: {str(e)}")
# Return a no-op meter provider
return MeterProvider(resource=resource)
def initialize_telemetry(
tracing_config: Optional[TracingConfig] = None,
metrics_config: Optional[MetricsConfig] = None,
) -> Tracer:
"""
Initialize OpenTelemetry with tracing and metrics.
Args:
tracing_config: Configuration for tracing
metrics_config: Configuration for metrics
Returns:
Tracer: Configured OpenTelemetry tracer
Example:
```python
config = TracingConfig(
service_name="amq-adapter",
service_version="1.0.0",
otlp_endpoint="http://jaeger:4317",
debug=True
)
metrics_config = MetricsConfig(
prometheus_port=9464,
prometheus_host="0.0.0.0"
)
tracer = initialize_telemetry(config, metrics_config)
```
"""
# Use default configs if none provided
if tracing_config is None:
tracing_config = TracingConfig(
service_name=os.getenv("SERVICE_NAME", "amq-adapter"),
service_version=os.getenv("SERVICE_VERSION", "unknown"),
otlp_endpoint=os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317"),
insecure=os.getenv("OTEL_INSECURE_MODE", "true").lower() == "true",
debug=os.getenv("OTEL_DEBUG", "false").lower() == "true",
)
if metrics_config is None:
metrics_config = MetricsConfig(
prometheus_port=int(os.getenv("PROMETHEUS_PORT", "9464")),
prometheus_host=os.getenv("PROMETHEUS_HOST", "localhost"),
export_interval_millis=int(os.getenv("METRICS_EXPORT_INTERVAL", "30000")),
)
try:
# Create resource
resource = create_resource(tracing_config)
# Set up tracing
provider = setup_trace_provider(tracing_config, resource)
trace.set_tracer_provider(provider)
logger.info(f"Initialized tracing for service {tracing_config.service_name}")
# Set up metrics
metrics_provider = setup_metrics_provider(metrics_config, resource)
metrics.set_meter_provider(metrics_provider)
logger.info("Initialized metrics export")
# Return configured tracer
return trace.get_tracer(
tracing_config.service_name,
tracing_config.service_version,
)
except Exception as e:
logger.error(f"Failed to initialize telemetry: {str(e)}")
# Return a no-op tracer in case of failure
return trace.get_tracer("no-op")
def create_span(tracer: Tracer, name: str, attributes: Optional[Dict[str, Any]] = None):
"""
Create a new span with the given name and attributes.
Args:
tracer: The OpenTelemetry tracer
name: Name of the span
attributes: Optional attributes to add to the span
Returns:
A context manager that creates and manages a span
"""
return tracer.start_as_current_span(
name,
attributes=attributes or {},
)
pyproject.toml
+6 -2
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "amq_adapter"
version = "0.2.31"
version = "0.2.32"
description = "The Python implementation of the AMQ Adapter for CleverMicro"
readme = "README.md"
authors = [
@@ -39,6 +39,9 @@ include = ["amqp*"]
[tool.black]
line-length = 100
[tool.pytest.ini_options]
asyncio_default_fixture_loop_scope = "module"
[project.optional-dependencies]
dev = [
"pytest>=7.0",
@@ -49,5 +52,6 @@ dev = [
"black",
"isort",
"flake8",
"pre-commit"
"pre-commit",
"pylint"
]
tests/__init__.py
View File
tests/adapter/__init__.py
View File
tests/adapter/test_backpressure_handler.py
+512
View File
@@ -0,0 +1,512 @@
#
# NOTE: need to install: pytest-asyncio
#
import asyncio
import json
from asyncio import AbstractEventLoop
from threading import Thread
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from aio_pika import Exchange
from aio_pika.abc import AbstractRobustChannel
from amqp.adapter.backpressure_handler import (
BackpressureHandler,
RunningAverage,
ScaleRequestV1,
)
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import ScalingRequestAlert
class TestScaleRequestV1:
"""
A class containing pytest unit tests for the ScaleRequestV1 class.
"""
def test_scale_request_v1_initialization(self):
"""
Test that the ScaleRequestV1 object is initialized correctly.
"""
service_id = "test-service"
task_id = "test-task"
max_availability = 10
current_availability = 5
request_type = ScalingRequestAlert.UPDATE
scale_request = ScaleRequestV1(
service_id,
task_id,
max_availability,
current_availability,
request_type,
)
assert scale_request.version == 1
assert scale_request.serviceId == service_id
assert scale_request.taskId == task_id
assert scale_request.maxAvailability == max_availability
assert scale_request.currentAvailability == current_availability
assert scale_request.requestType == request_type
assert scale_request.thread is None
def test_scale_request_v1_to_json(self):
"""
Test the to_json method of ScaleRequestV1.
"""
service_id = "test-service"
task_id = "test-task"
max_availability = 10
current_availability = 5
request_type = ScalingRequestAlert.UPDATE
scale_request = ScaleRequestV1(
service_id,
task_id,
max_availability,
current_availability,
request_type,
)
json_output = scale_request.to_json()
expected_json = json.dumps(
{
"version": 1,
"serviceId": service_id,
"taskId": task_id,
"maxAvailability": max_availability,
"currentAvailability": current_availability,
"requestType": request_type.value,
},
indent=2,
)
assert json_output == expected_json
def test_scale_request_v1_to_json_different_alert_type(self):
"""
Test the to_json method with a different ScalingRequestAlert type.
"""
service_id = "test-service"
task_id = "test-task"
max_availability = 10
current_availability = 5
request_type = ScalingRequestAlert.OVERLOAD # Change the alert type
scale_request = ScaleRequestV1(
service_id,
task_id,
max_availability,
current_availability,
request_type,
)
json_output = scale_request.to_json()
expected_json = json.dumps(
{
"version": 1,
"serviceId": service_id,
"taskId": task_id,
"maxAvailability": max_availability,
"currentAvailability": current_availability,
"requestType": request_type.value, # Ensure the value is correctly serialized
},
indent=2,
)
assert json_output == expected_json
class TestRunningAverage:
"""
A class containing pytest unit tests for the RunningAverage class.
"""
def test_running_average_initialization(self):
"""
Test that the RunningAverage object is initialized correctly.
"""
num_items = 5
ra = RunningAverage(num_items)
assert ra.buffer == [0] * num_items
assert ra.pointer == 0
assert ra.num_items == num_items
assert ra.is_filled is False
def test_running_average_add_one_value(self):
"""
Test adding a single value to the RunningAverage.
"""
ra = RunningAverage(3)
ra.add(10)
assert ra.buffer == [10, 0, 0]
assert ra.pointer == 1
assert ra.is_filled is False
def test_running_average_add_multiple_values_within_capacity(self):
"""
Test adding multiple values within the capacity of the RunningAverage buffer.
"""
ra = RunningAverage(3)
ra.add(5)
ra.add(10)
ra.add(15)
assert ra.buffer == [5, 10, 15]
assert ra.pointer == 0
assert ra.is_filled is True
def test_running_average_add_more_values_than_capacity(self):
"""
Test adding more values than the capacity of the RunningAverage buffer, checking wrap-around.
"""
ra = RunningAverage(3)
ra.add(1)
ra.add(2)
ra.add(3)
ra.add(4) # Overwrite the first value
ra.add(5) # Overwrite the second value
assert ra.buffer == [4, 5, 3]
assert ra.pointer == 2
assert ra.is_filled is True
def test_running_average_get_average_empty(self):
"""
Test get_average when no values have been added.
"""
ra = RunningAverage(3)
assert ra.get_average() == 0.0
def test_running_average_get_average_one_value(self):
"""
Test get_average after adding one value.
"""
ra = RunningAverage(3)
ra.add(10)
assert ra.get_average() == 10.0
def test_running_average_get_average_multiple_values_within_capacity(self):
"""
Test get_average with multiple values added within the buffer's capacity.
"""
ra = RunningAverage(3)
ra.add(1)
ra.add(2)
ra.add(3)
assert ra.get_average() == 2.0
def test_running_average_get_average_more_values_than_capacity(self):
"""
Test get_average after adding more values than the buffer's capacity (check wrap-around).
"""
ra = RunningAverage(3)
ra.add(1)
ra.add(2)
ra.add(3)
ra.add(4)
ra.add(5)
assert ra.get_average() == 4.0
def test_get_current_values_empty(self):
"""Test get_current_values when no values have been added."""
ra = RunningAverage(3)
assert ra.get_current_values() == []
def test_get_current_values_one_value(self):
"""Test get_current_values after adding one value."""
ra = RunningAverage(3)
ra.add(10)
assert ra.get_current_values() == [10]
def test_get_current_values_multiple_values_within_capacity(self):
"""Test get_current_values with multiple values within capacity."""
ra = RunningAverage(3)
ra.add(1)
ra.add(2)
ra.add(3)
assert ra.get_current_values() == [1, 2, 3]
def test_get_current_values_more_values_than_capacity(self):
"""Test get_current_values after adding more values than capacity (wrap-around)."""
ra = RunningAverage(3)
ra.add(1)
ra.add(2)
ra.add(3)
ra.add(4)
ra.add(5)
assert ra.get_current_values() == [3, 4, 5]
def test_get_current_values_wrap_around(self):
"""Test get_current_values when the buffer has wrapped around."""
ra = RunningAverage(4)
ra.add(1)
ra.add(2)
ra.add(3)
ra.add(4)
ra.add(5) # Overwrite 1
ra.add(6) # Overwrite 2
assert ra.get_current_values() == [3, 4, 5, 6]
# Test for BackpressureHandler
class TestBackpressureHandler:
"""
A class containing pytest unit tests for the BackpressureHandler class.
"""
@pytest.fixture
def mock_channel(self):
"""
Pytest fixture to create a mock AbstractRobustChannel.
"""
return AsyncMock(spec=AbstractRobustChannel)
@pytest.fixture
def mock_loop(self):
"""
Pytest fixture to create a mock AbstractEventLoop.
"""
return MagicMock(spec=AbstractEventLoop)
@pytest.fixture
def mock_config(self):
"""
Pytest fixture to create a mock AMQConfiguration.
"""
return AMQConfiguration("")
def test_backpressure_handler_initialization(self, mock_channel, mock_loop, mock_config):
"""
Test that the BackpressureHandler object is initialized correctly.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
assert handler.channel == mock_channel
assert handler.loop == mock_loop
assert handler.config == mock_config
assert handler.swarm_service_id == "clever-amqp-service"
assert handler.swarm_task_id == "python-adapter"
assert handler.parallel_workers == 1
assert handler.average_cpu_usage is None
def test_increase_parallel_executions(self, mock_channel, mock_loop, mock_config):
"""
Test the increase_parallel_executions method.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
initial_executions = handler.current_parallel_executions
handler.increase_parallel_executions()
assert handler.current_parallel_executions == initial_executions + 1
def test_decrease_parallel_executions(self, mock_channel, mock_loop, mock_config):
"""
Test the decrease_parallel_executions method.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
handler.current_parallel_executions = 5
handler.decrease_parallel_executions()
assert handler.current_parallel_executions == 4
handler.current_parallel_executions = 0
handler.decrease_parallel_executions()
assert handler.current_parallel_executions == 0 # Should not go below 0
def test_update_parallel_executions(self, mock_channel, mock_loop, mock_config):
"""
Test the update_parallel_executions method.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
handler.update_parallel_executions(100)
assert handler.current_parallel_executions == 100
def test_update_last_data_message_time(self, mock_channel, mock_loop, mock_config):
"""
Test the update_last_data_message_time method.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
with patch("time.time") as mock_time:
mock_time.return_value = 12345
handler.update_last_data_message_time()
assert handler.last_data_message_time == 12345
def test_start_backpressure_monitor(self, mock_channel, mock_loop, mock_config):
"""
Test the start_backpressure_monitor method.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
handler.do_loop = 1 # allow only 1 execution of the loop and then exit the thread, else tests will never end
thread = handler.start_backpressure_monitor()
assert isinstance(thread, Thread)
assert handler.thread == thread
@pytest.mark.asyncio
async def test_check_overload_condition_no_overload(self, mock_channel, mock_loop, mock_config):
"""
Test the check_overload_condition method when there is no overload.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
handler.current_parallel_executions = 1
handler.parallel_workers = 5
handler.last_backpressure_event = ScalingRequestAlert.UPDATE
with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload:
await handler.check_overload_condition()
mock_handle_overload.assert_not_called()
@pytest.mark.asyncio
async def test_check_overload_condition_overload(self, mock_channel, mock_loop, mock_config):
"""
Test the check_overload_condition method when there is an overload.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
handler.current_parallel_executions = 9 # Assuming parallel_workers is 10
handler.last_backpressure_event = ScalingRequestAlert.UPDATE
with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload:
with patch("time.time") as mock_time:
mock_time.return_value = 12345
await handler.check_overload_condition()
mock_handle_overload.assert_called_once()
assert handler.last_backpressure_event_time == 12345
assert handler.last_backpressure_event == ScalingRequestAlert.OVERLOAD
@pytest.mark.asyncio
async def test_check_overload_condition_already_overload(
self, mock_channel, mock_loop, mock_config
):
"""
Test the check_overload_condition method when the last event was already an overload.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
handler.current_parallel_executions = 10
handler.last_backpressure_event = ScalingRequestAlert.OVERLOAD
with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload:
await handler.check_overload_condition()
mock_handle_overload.assert_not_called()
# backpressure_monitor_loop is difficult to test directly without significant mocking
# and potentially long-running tests. The following is a basic outline of how
# you might approach testing it, focusing on verifying that the correct methods
# are called under specific conditions. This is NOT a complete, runnable test.
@pytest.mark.asyncio
async def test_backpressure_monitor_loop(self, mock_channel, mock_loop, mock_config):
"""
Test the backpressure_monitor_loop method (outline).
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
handler.average_cpu_usage = RunningAverage(
100, initial_value=99
) # Mock the RunningAverage in OVERLOAD state
handler.average_cpu_usage.is_filled = True # Together with 0 values ensures OVERLOAD state
# Setup mocks for dependent methods
handler.handle_backpressure_overload_event = AsyncMock()
handler._handle_backpressure_idle_event = AsyncMock()
handler.publish_backpressure_request = AsyncMock()
# 1. Simulate an overload condition:
handler.last_backpressure_event_time = 0 # Force event trigger
handler.last_data_message_time = 0
handler.do_loop = 1
handler.config.backpressure.cpu_overload_duration = (
-1
) # Force overload condition even for recent change
await handler._backpressure_monitor()
handler.handle_backpressure_overload_event.assert_called_once()
# 2. Simulate an idle condition
handler.average_cpu_usage = RunningAverage(
100, initial_value=0
) # Mock the RunningAverage in IDLE state
handler.average_cpu_usage.is_filled = True # Together with 0 values ensures IDLE state
handler.last_backpressure_event_time = 0
handler.last_data_message_time = 0
handler.do_loop = 1
handler.config.backpressure.idle_duration = (
-1
) # Force idle condition even for recent change
await handler._backpressure_monitor()
handler._handle_backpressure_idle_event.assert_called_once()
# 3. Simulate neither idle nor overload, but time_window passed
handler.handle_backpressure_overload_event.call_count = 0
handler._handle_backpressure_idle_event.call_count = 0
handler.average_cpu_usage = RunningAverage(100) # Mock the RunningAverage in pritine state
handler.last_backpressure_event_time = 0
handler.last_data_message_time = 0
handler.do_loop = 1
await handler._backpressure_monitor()
assert (
handler.publish_backpressure_request.call_count
+ handler.handle_backpressure_overload_event.call_count
+ handler._handle_backpressure_idle_event.call_count
) > 0
@pytest.mark.asyncio
async def test_handle_backpressure_overload_event(self, mock_channel, mock_loop, mock_config):
"""
Test the handle_backpressure_overload_event method.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
handler.publish_backpressure_request = AsyncMock()
with patch("time.time") as mock_time:
mock_time.return_value = 12345
await handler.handle_backpressure_overload_event()
assert handler.publish_backpressure_request.call_count == 1
assert handler.last_data_message_time == 12345
@pytest.mark.asyncio
async def test_handle_backpressure_idle_event(self, mock_channel, mock_loop, mock_config):
"""
Test the _handle_backpressure_idle_event method.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
handler.publish_backpressure_request = AsyncMock()
with patch("time.time") as mock_time:
mock_time.return_value = 12345
await handler._handle_backpressure_idle_event()
assert handler.publish_backpressure_request.call_count == 1
assert handler.last_data_message_time == 12345
@pytest.mark.asyncio
async def test_publish_backpressure_request_exchange_exists(
self, mock_channel, mock_loop, mock_config
):
"""
Test the publish_backpressure_request method when the exchange already exists.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
handler.exchange = AsyncMock(spec=Exchange)
with patch("asyncio.run_coroutine_threadsafe") as mock_async_run:
mock_async_run.return_value = asyncio.Future()
mock_async_run.return_value.set_result(None)
await handler.publish_backpressure_request(
ScaleRequestV1("test_service", "test_task", 100, 50, ScalingRequestAlert.UPDATE)
)
assert mock_async_run.call_count == 1
# handler.exchange.publish.assert_called_once()
@pytest.mark.asyncio
async def test_publish_backpressure_request_exchange_does_not_exist(
self, mock_channel, mock_loop, mock_config
):
"""
Test the publish_backpressure_request method when the exchange does not exist and needs to be created.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
handler.exchange = None # Simulate exchange not existing
mock_exchange = AsyncMock(spec=Exchange)
mock_channel.get_exchange.return_value = mock_exchange
BackpressureHandler._callback_list = {}
with patch("asyncio.run_coroutine_threadsafe") as mock_async_run:
mock_async_run.return_value = asyncio.Future()
mock_async_run.return_value.set_result(
mock_exchange
) # Make it return the mock exchange
await handler.publish_backpressure_request(
ScaleRequestV1("test_service", "test_task", 100, 50, ScalingRequestAlert.UPDATE)
)
assert mock_async_run.call_count == 2
assert "_wrap_rabbit_mq_api_init" in BackpressureHandler._callback_list
assert "_wrap_rabbit_mq_api" in BackpressureHandler._callback_list
# Since I need also to mockup the asyncio.run_coroutine_threadsafe, the following mockups are
# not called because mockup overrides that call with preset return value.
# mock_channel.get_exchange.assert_called_once_with(name="cleverthis.clevermicro.management")
# handler.exchange.publish.assert_called_once()
tests/adapter/test_cleverthis_service_adapter.py
+123 -61
View File
@@ -1,14 +1,82 @@
import asyncio
import unittest
from unittest.mock import AsyncMock, MagicMock, patch
from aio_pika.abc import AbstractChannel, AbstractRobustConnection
from aiohttp import ClientResponse
from aiohttp.web_fileresponse import FileResponse
from amqp.adapter.cleverthis_service_adapter import AMQMessage, CleverThisServiceAdapter
from amqp.adapter.cleverthis_service_adapter import (
AMQMessage,
CleverThisServiceAdapter,
_convert_file_response,
_create_reply_exchange_and_queue,
is_likely_json,
)
from amqp.adapter.file_handler import FileHandler
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import DataMessage, DataResponse
from amqp.model.model import DataResponse
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE
class TestCleverThisServiceAdapter(unittest.TestCase):
class TestCleverThisServiceAdapter(unittest.IsolatedAsyncioTestCase):
async def test_create_reply_exchange_and_queue_success(self):
connection_mock = AsyncMock()
channel_mock = AsyncMock()
exchange_mock = AsyncMock()
queue_mock = AsyncMock()
connection_mock.channel.return_value = channel_mock
channel_mock.get_exchange.return_value = exchange_mock
channel_mock.declare_queue.return_value = queue_mock
queue_mock.name = "test_queue_name"
channel, exchange, queue_name = await _create_reply_exchange_and_queue(connection_mock)
assert channel == channel_mock
assert exchange == exchange_mock
assert queue_name == "test_queue_name"
channel_mock.get_exchange.assert_called_once_with(name=AMQ_REPLY_TO_EXCHANGE, ensure=True)
channel_mock.declare_queue.assert_called_once_with(exclusive=False)
queue_mock.bind.assert_called_once_with(
exchange=AMQ_REPLY_TO_EXCHANGE, routing_key="test_queue_name"
)
async def test_convert_file_response_successful(self):
connection_mock = AsyncMock(spec=AbstractRobustConnection)
channel_mock: AbstractChannel = AsyncMock(spec=AbstractChannel)
exchange_mock = AsyncMock()
queue_mock = AsyncMock()
connection_mock.channel = AsyncMock(return_value=channel_mock)
channel_mock.get_exchange.return_value = exchange_mock
channel_mock.declare_queue.return_value = queue_mock
channel_mock.close = AsyncMock()
queue_mock.name = "test_queue_name"
mock_file_handler = MagicMock(spec=FileHandler)
mock_file_handler.publish_file = AsyncMock(return_value=("test_stream", 12345))
mock_file_response = MagicMock(spec=FileResponse)
mock_file_response.path = "/test/path"
mock_file_response.filename = "test_file.txt"
mock_loop = asyncio.get_event_loop()
json_response = await _convert_file_response(
mock_file_response, connection_mock, mock_file_handler, mock_loop
)
assert '"filename": "test_file.txt"' in json_response
assert '"streamName": "test_stream"' in json_response
assert '"fileSize": 12345' in json_response
def test_is_likely_json(self):
assert is_likely_json("{ }") is True
assert is_likely_json("[ ]") is True
assert is_likely_json("{]") is False
assert is_likely_json("[}") is False
assert is_likely_json("something else") is False
def setUp(self):
# Mock AMQConfiguration
self.mock_amq_config = AMQConfiguration("./application.properties.local")
@@ -20,9 +88,10 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
self.adapter = CleverThisServiceAdapter(
amq_configuration=self.mock_amq_config, routes=[], session=self.mock_session
)
self.adapter.service_host = "testhost"
def test_init(self):
self.assertEqual(self.adapter.service_host, "localhost")
self.assertEqual(self.adapter.service_host, "testhost")
self.assertEqual(self.adapter.service_port, 8080)
self.assertFalse(self.adapter.require_authenticated_user)
@@ -48,26 +117,27 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
@patch.object(CleverThisServiceAdapter, "get_current_user")
async def test_on_message_unauthenticated(self, mock_get_user):
mock_message = MagicMock(spec=AMQMessage)
mock_message.headers = {}
mock_message.path = "/some/path"
mock_message.method = "GET"
mock_message.id = "123"
mock_message.trace_info = {}
mock_message.headers.return_value = {}
mock_message.path.return_value = "/some/path"
mock_message.method.return_value = "GET"
mock_message.id.return_value = "123"
mock_message.trace_info.return_value = {}
# Test when authentication is required but not provided
self.adapter.require_authenticated_user = True
response = await self.adapter.on_message(mock_message)
self.assertEqual(response.response_code(), "401")
mock_get_user.assert_not_called()
@patch.object(CleverThisServiceAdapter, "get_current_user")
@patch.object(CleverThisServiceAdapter, "get")
async def test_on_message_authenticated(self, mock_get, mock_get_user):
async def test_on_message_authenticated(self, mock_get: MagicMock, mock_get_user):
mock_message = MagicMock(spec=AMQMessage)
mock_message.headers = {"Authorization": ["Bearer valid_token"]}
mock_message.path = "/some/path"
mock_message.method = "GET"
mock_message.id = "123"
mock_message.trace_info = {}
mock_message.headers.return_value = {"Authorization": ["Bearer valid_token"]}
mock_message.path.return_value = "/some/path"
mock_message.method.return_value = "GET"
mock_message.id.return_value = "123"
mock_message.trace_info.return_value = {}
mock_user = MagicMock()
mock_get_user.return_value = mock_user
@@ -75,6 +145,7 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
mock_response = MagicMock(spec=DataResponse)
mock_get.return_value = mock_response
self.adapter.require_authenticated_user = True
response = await self.adapter.on_message(mock_message)
self.assertEqual(response, mock_response)
@@ -84,11 +155,11 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
@patch.object(CleverThisServiceAdapter, "get_current_user")
async def test_on_message_login_path(self, mock_get_user):
mock_message = MagicMock(spec=AMQMessage)
mock_message.headers = {}
mock_message.path = "/login"
mock_message.method = "POST"
mock_message.id = "123"
mock_message.trace_info = {}
mock_message.headers.return_value = {}
mock_message.path.return_value = "/login"
mock_message.method.return_value = "POST"
mock_message.id.return_value = "123"
mock_message.trace_info.return_value = {}
mock_response = MagicMock(spec=DataResponse)
with patch.object(self.adapter, "post", return_value=mock_response) as mock_post:
@@ -108,11 +179,11 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
for method in methods:
with self.subTest(method=method):
mock_message = MagicMock(spec=AMQMessage)
mock_message.headers = {"Authorization": ["Bearer valid_token"]}
mock_message.path = "/some/path"
mock_message.method = method
mock_message.id = "123"
mock_message.trace_info = {}
mock_message.headers.return_value = {"Authorization": ["Bearer valid_token"]}
mock_message.path.return_value = "/some/path"
mock_message.method.return_value = method
mock_message.id.return_value = "123"
mock_message.trace_info.return_value = {}
with patch.object(
self.adapter, method.lower(), return_value=mock_response
@@ -122,11 +193,11 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
mock_method.assert_called_once()
async def test_get_with_session(self):
mock_message = MagicMock(spec=DataMessage)
mock_message.path = "/api/test"
mock_message.headers = {"X-Test": "value"}
mock_message.trace_info = {"trace-id": "123"}
mock_message.id = "msg123"
mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.path.return_value = "/api/test"
mock_message.headers.return_value = {"X-Test": "value"}
mock_message.trace_info.return_value = {"trace-id": "123"}
mock_message.id.return_value = "msg123"
mock_response = MagicMock(spec=ClientResponse)
mock_response.status = 200
@@ -146,26 +217,26 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
async def test_get_without_session(self):
self.adapter.session = None
mock_message = MagicMock(spec=DataMessage)
mock_message.path = "/api/test"
mock_message.id = "msg123"
mock_message.trace_info = {}
mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.path.return_value = "/api/test"
mock_message.id.return_value = "msg123"
mock_message.trace_info.return_value = {}
response = await self.adapter.get(mock_message, None)
self.assertEqual(response.response_code(), "404")
async def test_put_with_form_handling(self):
mock_message = MagicMock(spec=DataMessage)
mock_message.path = "/api/test"
mock_message.headers = {}
mock_message.trace_info = {}
mock_message.id = "msg123"
mock_message.body = b"test=value"
mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.path.return_value = "/api/test"
mock_message.headers.return_value = {}
mock_message.trace_info.return_value = {}
mock_message.id.return_value = "msg123"
mock_message.body.return_value = b"test=value"
mock_response = MagicMock(spec=ClientResponse)
mock_response.status = 200
mock_response.content_type = "application/json"
mock_response: DataResponse = MagicMock(spec=DataResponse)
mock_response.response_code.return_value = "200"
mock_response.content_type.return_value = "application/json"
mock_response.read = AsyncMock(return_value=b'{"result": "success"}')
self.mock_session.put.return_value = mock_response
@@ -181,15 +252,15 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
mock_handler.assert_called_once()
async def test_post_with_form_handling(self):
mock_message = MagicMock(spec=DataMessage)
mock_message.path = "/api/test"
mock_message.headers = {}
mock_message.trace_info = {}
mock_message.id = "msg123"
mock_message.body = b"test=value"
mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.path.return_value = "/api/test"
mock_message.headers.return_value = {}
mock_message.trace_info.return_value = {}
mock_message.id.return_value = "msg123"
mock_message.body.return_value = b"test=value"
mock_response = MagicMock(spec=ClientResponse)
mock_response.status = 201
mock_response: DataResponse = MagicMock(spec=DataResponse)
mock_response.response_code.return_value = "201"
mock_response.content_type = "application/json"
mock_response.read = AsyncMock(return_value=b'{"result": "created"}')
@@ -201,25 +272,16 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
wraps=self.adapter._authorize_forward_request,
) as mock_handler:
response = await self.adapter.post(mock_message, None)
self.assertEqual(response.response_code(), "201")
mock_handler.assert_called_once()
async def test_head_not_implemented(self):
mock_message = MagicMock(spec=DataMessage)
mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.id = "msg123"
with self.assertRaises(NotImplementedError):
await self.adapter.head(mock_message, None)
async def test_delete_not_implemented(self):
mock_message = MagicMock(spec=DataMessage)
mock_message.id = "msg123"
# Assuming delete is not implemented in base class
with self.assertRaises(NotImplementedError):
await self.adapter.delete(mock_message, None)
if __name__ == "__main__":
unittest.main()
tests/adapter/test_data_parser.py
+165
View File
@@ -0,0 +1,165 @@
import json
import unittest
from unittest.mock import MagicMock, patch
from amqp.adapter.data_parser import AMQDataParser, MultipartFormDataParser
class DummyAMQMessage:
def __init__(self, headers=None, body=b"", extra_data=None):
self._headers = headers or {"Content-Type": ["application/json"]}
self._body = body
self.extra_data = extra_data if extra_data is not None else {}
def headers(self):
return self._headers
def body(self):
return self._body
class TestAMQDataParser(unittest.TestCase):
def test_process_form_data(self):
cases = [
(
'{"form": {"foo": ["bar"]}, "files": {"file1": ["data"]}}',
{"foo": "bar", "file1": "data"},
),
({"form": {"a": ["1"]}, "files": {"b": ["2"]}}, {"a": "1", "b": "2"}),
({"form": {}, "files": {}}, {}),
]
for json_input, expected in cases:
result = AMQDataParser.process_form_data(json_input)
self.assertEqual(result, expected)
def test_parse_get_url(self):
url = "http://example.com/path?foo=1&bar=2&foo=3"
path, params = AMQDataParser.parse_get_url(url)
self.assertEqual(path, "/path")
self.assertTrue(params["foo"] == ["1", "3"] or params["foo"] == "1" or params["foo"] == "3")
self.assertEqual(params["bar"], "2")
def test_parse_request_body_json(self):
data = {"foo": "bar"}
msg = DummyAMQMessage(
headers={"Content-Type": ["application/json"]}, body=json.dumps(data).encode()
)
result = AMQDataParser.parse_request_body(msg)
self.assertEqual(result, data)
def test_parse_request_body_invalid_json(self):
msg = DummyAMQMessage(headers={"Content-Type": ["application/json"]}, body=b"notjson")
with patch("amqp.adapter.data_parser.logging_error") as log_err:
with self.assertRaises(ValueError):
AMQDataParser.parse_request_body(msg)
self.assertTrue(log_err.called)
def test_parse_request_body_urlencoded(self):
msg = DummyAMQMessage(
headers={"Content-Type": ["application/x-www-form-urlencoded"]}, body=b"foo=bar&baz=qux"
)
result = AMQDataParser.parse_request_body(msg)
self.assertEqual(result, {b"foo": b"bar", b"baz": b"qux"})
def test_parse_request_body_urlencoded_json(self):
data = {"form": {"foo": ["bar"]}, "files": {}}
msg = DummyAMQMessage(
headers={"Content-Type": ["application/x-www-form-urlencoded"]},
body=json.dumps(data).encode(),
)
result = AMQDataParser.parse_request_body(msg)
self.assertEqual(result, {"foo": "bar"})
def test_parse_request_body_invalid_urlencoded(self):
msg = DummyAMQMessage(
headers={"Content-Type": ["application/x-www-form-urlencoded"]}, body=b"not a json"
)
with patch("amqp.adapter.data_parser.logging_error") as log_err:
with self.assertRaises(ValueError):
AMQDataParser.parse_request_body(msg)
self.assertTrue(log_err.called)
def test_parse_request_body_text(self):
msg = DummyAMQMessage(headers={"Content-Type": ["text/plain"]}, body=b"hello world")
result = AMQDataParser.parse_request_body(msg)
self.assertEqual(result, {"text": b"hello world"})
def test_parse_request_body_xml(self):
xml = b"<root><foo>bar</foo><baz>qux</baz></root>"
msg = DummyAMQMessage(headers={"Content-Type": ["application/xml"]}, body=xml)
result = AMQDataParser.parse_request_body(msg)
self.assertEqual(result, {"foo": "bar", "baz": "qux"})
def test_parse_request_body_invalid_xml(self):
msg = DummyAMQMessage(
headers={"Content-Type": ["application/xml"]}, body=b"<root><foo></bar></root>"
)
with self.assertRaises(ValueError):
AMQDataParser.parse_request_body(msg)
def test_parse_request_body_unsupported(self):
msg = DummyAMQMessage(headers={"Content-Type": ["application/unknown"]}, body=b"")
with self.assertRaises(ValueError):
AMQDataParser.parse_request_body(msg)
class TestMultipartFormDataParser(unittest.TestCase):
def test_multipart_form_data_parser_json(self):
data = {"form": {"foo": ["bar"]}, "files": {}}
msg = DummyAMQMessage(
headers={"Content-Type": ["multipart/form-data"]},
body=json.dumps(data).encode(),
extra_data={"some": "data"},
)
parser = MultipartFormDataParser(msg)
self.assertTrue(parser.is_json)
self.assertEqual(parser.form_data, {"foo": "bar"})
self.assertTrue(parser.is_valid)
def test_multipart_form_data_parser_invalid_json(self):
msg = DummyAMQMessage(
headers={"Content-Type": ["multipart/form-data"]},
body=b"notjson",
extra_data={"some": "data"},
)
with patch("amqp.adapter.data_parser.logging_error") as log_err:
parser = MultipartFormDataParser(msg)
self.assertFalse(parser.is_valid)
self.assertTrue(log_err.called)
def test_multipart_form_data_parser_fallback_to_json(self):
msg = DummyAMQMessage(
headers={"Content-Type": ["multipart/form-data"]},
body=json.dumps({"form": {"foo": ["bar"]}, "files": {}}).encode(),
)
with patch(
"amqp.adapter.data_parser.python_multipart.parse_form", side_effect=Exception("fail")
):
parser = MultipartFormDataParser(msg)
self.assertTrue(parser.is_json)
self.assertEqual(parser.form_data, {"foo": "bar"})
self.assertTrue(parser.is_valid)
def test_multipart_form_data_parser_on_field_and_file(self):
msg = DummyAMQMessage(
headers={"Content-Type": ["multipart/form-data"]}, body=b"", extra_data={}
)
parser = MultipartFormDataParser(msg)
class DummyField:
field_name = b"foo"
value = b"bar"
class DummyFile:
field_name = b"file"
file_object = MagicMock()
file_object.seek = MagicMock()
size = 123
file_name = b"filename.txt"
parser.on_field(DummyField())
self.assertEqual(parser.form_data["foo"], b"bar")
parser.on_file(DummyFile())
self.assertIn("file", parser.form_data)
self.assertTrue(hasattr(parser.form_data["file"], "filename"))
self.assertEqual(parser.form_data["file"].filename, "filename.txt")
tests/adapter/test_file_handler.py
+263 -2
View File
@@ -1,11 +1,22 @@
import asyncio
import os
import tempfile
from unittest import TestCase
from asyncio import AbstractEventLoop
from collections import defaultdict
from pathlib import Path
from shutil import rmtree
from threading import Thread
from unittest import IsolatedAsyncioTestCase
from unittest.mock import AsyncMock, MagicMock, mock_open, patch
from aio_pika import IncomingMessage
from aio_pika.abc import AbstractRobustConnection
from amqp.adapter import file_handler
from amqp.adapter.file_handler import StreamingFileHandler
class TestFileHandler(TestCase):
class TestFileHandler(IsolatedAsyncioTestCase):
def test_ensure_directory_exists(self):
"""
Test function for ensure_directory_exists. Creates a temporary directory
@@ -48,3 +59,253 @@ class TestFileHandler(TestCase):
), f"Test Case 3 Version 3 Failed: {result3_v3} != {os.path.join(temp_dir, 'test3.3.txt')}"
print("All test cases passed!")
async def test_download_from_queue_to_file(self):
loop = asyncio.new_event_loop()
tempFile = tempfile.mktemp()
file_handler = StreamingFileHandler()
with patch.object(
StreamingFileHandler, "_get_connection", new_callable=AsyncMock
) as mock_connect:
mock_connection: AbstractRobustConnection = MagicMock(spec=AbstractRobustConnection)
mock_connect.return_value = mock_connection
mock_connection.channel = AsyncMock()
mock_channel = mock_connection.channel.return_value
mock_channel.get_queue = AsyncMock()
mock_queue = mock_channel.get_queue.return_value
mock_queue.iterator = MagicMock()
mock_queue_iter = mock_queue.iterator().__aenter__.return_value
mock_message: IncomingMessage = MagicMock(spec=IncomingMessage)
mock_message.body = b"test_data"
mock_message.body_size = len(b"test_data")
mock_message.headers = {"eof": 1}
mock_queue_iter.__aiter__.return_value = [mock_message]
# Act
result = await file_handler._download_from_queue_to_file(
"test_rabbitmq",
"test_queue",
filepath=tempFile,
loop=loop,
)
# clean up
os.remove(tempFile)
# Assert
assert result == (len(b"test_data"), 1)
mock_connect.assert_called_once_with("test_rabbitmq", loop=loop)
mock_channel.get_queue.assert_called_once_with("test_queue", ensure=False)
async def test_download_file_success(self):
# Arrange
file_handler = StreamingFileHandler()
loop = asyncio.new_event_loop()
loop_thread = Thread(target=loop.run_forever, daemon=True)
loop_thread.start()
rabbitmq_url = "amqp://test.rabbitmq"
file_info = {
"filename": "test_file.txt",
"size": 1024,
"streamName": "test_stream",
"queue_name": "test_queue",
"key": "test_key",
}
message_data = {}
output_dir = tempfile.mkdtemp()
file_handler.ensure_directory_exists = MagicMock(return_value=f"{output_dir}/test_file.txt")
file_handler._download_from_queue_to_file = AsyncMock(return_value=(123, 1))
# Act
result = await file_handler.download_file(
loop, rabbitmq_url, file_info, message_data, output_dir
)
# clean up
loop.stop()
rmtree(output_dir, ignore_errors=True)
# Assert
assert result == (123, 1)
file_handler.ensure_directory_exists.assert_called_once_with(f"{output_dir}/test_file.txt")
file_handler._download_from_queue_to_file.assert_called_once_with(
"amqp://test.rabbitmq", "test_queue", f"{output_dir}/test_file.txt", loop
)
async def test_download_file_queue_not_found(self):
# Arrange
file_handler = StreamingFileHandler()
loop = MagicMock(spec=AbstractEventLoop)
rabbitmq_url = "amqp://test.rabbitmq"
file_info = {
"filename": "test_file.txt",
"size": 1024,
"streamName": "test_stream",
"queue_name": None,
"key": "test_key",
}
message_data = {}
output_dir = "test_output"
# Act
result = await file_handler.download_file(
loop, rabbitmq_url, file_info, message_data, output_dir
)
# Assert
assert result == (0, 2)
async def test_download_from_queue_to_buffer(self):
loop = asyncio.new_event_loop()
file_handler = StreamingFileHandler()
with patch.object(
StreamingFileHandler, "_get_connection", new_callable=AsyncMock
) as mock_connect:
mock_connection: AbstractRobustConnection = MagicMock(spec=AbstractRobustConnection)
mock_connect.return_value = mock_connection
mock_connection.channel = AsyncMock()
mock_channel = mock_connection.channel.return_value
mock_channel.get_queue = AsyncMock()
mock_queue = mock_channel.get_queue.return_value
mock_queue.iterator = MagicMock()
mock_queue_iter = mock_queue.iterator().__aenter__.return_value
mock_message: IncomingMessage = MagicMock(spec=IncomingMessage)
mock_message.body = b"test_data"
mock_message.body_size = len(b"test_data")
mock_message.headers = {"eof": 1}
mock_queue_iter.__aiter__.return_value = [mock_message]
# Act
result = await file_handler._download_from_queue_to_buffer(
"test_rabbitmq",
"test_queue",
loop=loop,
)
# Assert
assert result == (b"test_data", 1)
mock_connect.assert_called_once_with("test_rabbitmq", loop=loop)
mock_channel.get_queue.assert_called_once_with("test_queue", ensure=False)
async def test_download_buffer_success(self):
# Arrange
file_handler = StreamingFileHandler()
loop = asyncio.new_event_loop()
loop_thread = Thread(target=loop.run_forever, daemon=True)
loop_thread.start()
rabbitmq_url = "amqp://test.rabbitmq"
file_handler._download_from_queue_to_buffer = AsyncMock(return_value=(b"hello", 1))
# Act
result = await file_handler.download_buffer(loop, rabbitmq_url, "test_queue")
# clean up
loop.stop()
# Assert
assert result == (b"hello", 1)
file_handler._download_from_queue_to_buffer.assert_called_once_with(
"amqp://test.rabbitmq", "test_queue", loop
)
async def test_download_buffer_queue_not_found(self):
# Arrange
file_handler = StreamingFileHandler()
loop = MagicMock(spec=AbstractEventLoop)
rabbitmq_url = "amqp://test.rabbitmq"
# Act
result = await file_handler.download_buffer(loop, rabbitmq_url, None)
# Assert
assert result == (0, 2)
async def test_on_message_no_files_to_download(self):
handler = StreamingFileHandler()
message = MagicMock()
message.body = b"{}"
json_data = b'{"files": {}}'
rabbitmq_url = "amqp://localhost"
loop = AsyncMock()
output_dir = "downloaded_files"
handler.download_file = AsyncMock()
handler.download_file.return_value = (123, 1)
with patch("amqp.adapter.file_handler.logging_info"):
result = await handler.on_message(message, json_data, rabbitmq_url, loop, output_dir)
assert isinstance(result, defaultdict)
self.assertEqual(result, defaultdict())
async def test_on_message_with_files_to_download(self):
handler = StreamingFileHandler()
message = MagicMock()
message.body = b'{"fileStream1": "queue1"}'
json_data = b'{"files": {"key1": [{"streamName": "fileStream1"}]}}'
rabbitmq_url = "amqp://localhost"
loop = AsyncMock()
output_dir = "downloaded_files"
handler.download_file = AsyncMock()
with patch("amqp.adapter.file_handler.logging_info") as mock_logging_info:
result = await handler.on_message(message, json_data, rabbitmq_url, loop, output_dir)
self.assertEqual(result, file_handler.download_locations)
handler.download_file.assert_awaited_once()
mock_logging_info.assert_any_call(
f" [{message.delivery_tag}] about to create task for {{'streamName': 'fileStream1', 'queue_name': 'queue1', 'key': 'key1'}}"
)
mock_logging_info.assert_called_with(f"All files downloaded. d-tag: {message.delivery_tag}")
async def test_on_message_handles_exception(self):
handler = StreamingFileHandler()
message = MagicMock()
message.body = b"invalid_json"
message.ack = AsyncMock()
json_data = b"invalid_json"
rabbitmq_url = "amqp://localhost"
loop = AsyncMock()
output_dir = "downloaded_files"
with patch("amqp.adapter.file_handler.logging_error") as mock_logging_error:
await handler.on_message(message, json_data, rabbitmq_url, loop, output_dir)
mock_logging_error.assert_called()
async def test_publish_file_file_not_found(self):
file_handler = StreamingFileHandler()
loop = asyncio.get_event_loop()
exchange_mock = AsyncMock()
file_path = Path("/path/to/nonexistent_file.txt")
with self.assertRaises(FileNotFoundError):
await file_handler.publish_file(
exchange=exchange_mock,
file_path=file_path,
routing_key="test_routing_key",
loop=loop,
)
async def test_publish_file_success(self):
file_handler = StreamingFileHandler()
loop = asyncio.get_event_loop()
exchange_mock = AsyncMock()
file_path = Path("/path/to/test_file.txt")
data = b"test-data"
with (
patch("amqp.adapter.file_handler.Path.is_file", return_value=True),
patch("amqp.adapter.file_handler.os.path.getsize", return_value=len(data)),
patch("builtins.open", new=mock_open(read_data=data)) as mock_file,
):
result = await file_handler.publish_file(
exchange=exchange_mock,
file_path=file_path,
routing_key="test_routing_key",
loop=loop,
max_chunk_size=512,
)
mock_file.assert_called_with(file_path, "rb")
self.assertEqual(result, ("test_routing_key", len(data)))
tests/adapter/test_logging_utils.py
+138
View File
@@ -0,0 +1,138 @@
import os
from importlib.metadata import version
from unittest.mock import patch
from amqp.adapter.logging_utils import (
get_context_info,
logging_debug,
logging_error,
logging_info,
logging_warning,
)
__version__ = version("amq_adapter") # Name must match pyproject.toml [project] name
__verbose__ = os.environ.get("AMQ_VERBOSE_LOGGING", "0") == "1"
class TestLoggingFunctions:
"""
A class containing pytest unit tests for the logging functions.
"""
@patch("logging.error")
def test_logging_error_with_exception(self, mock_logger):
"""
Test logging_error function when called within an exception handler.
"""
try:
raise ValueError("Test error")
except ValueError as ve:
logging_error("An error occurred: %s", ve)
mock_logger.assert_called_once()
# Assert that the traceback information is included in the log message
assert len(mock_logger.call_args_list) == 1
_val = mock_logger.call_args_list[0][0][0]
assert "[E] 'An error occurred: %s'" in _val
assert "test_logging_error_with_exception" in _val
@patch("logging.error")
def test_logging_error_without_exception(self, mock_logger):
"""
Test logging_error function when called without an active exception.
"""
logging_error("A regular error")
mock_logger.assert_called_once()
assert "A regular error" in mock_logger.call_args_list[0][0][0]
@patch("logging.warning")
def test_logging_warning_verbose_on(self, mock_logger):
"""
Test logging_warning function when verbose logging is enabled.
"""
with patch(
"amqp.adapter.logging_utils.__verbose__", True
): # Simulate verbose logging
logging_warning("A warning message")
mock_logger.assert_called_once()
assert "A warning message" in mock_logger.call_args_list[0][0][0]
@patch("logging.warning")
def test_logging_warning_verbose_off(self, mock_logger):
"""
Test logging_warning function when verbose logging is disabled.
"""
with patch(
"amqp.adapter.logging_utils.__verbose__", False
): # Simulate non-verbose logging
logging_warning("A warning message")
mock_logger.assert_called_once()
assert "A warning message" in mock_logger.call_args_list[0][0][0]
@patch("logging.info")
def test_logging_info_verbose_on(self, mock_logger):
"""
Test logging_info function with verbose logging enabled.
"""
with patch("amqp.adapter.logging_utils.__verbose__", True):
logging_info("An info message")
mock_logger.assert_called_once()
assert "An info message" in mock_logger.call_args_list[0][0][0]
@patch("logging.info")
def test_logging_info_verbose_off(self, mock_logger):
"""
Test logging_info function with verbose logging disabled.
"""
with patch("amqp.adapter.logging_utils.__verbose__", False):
logging_info("An info message")
mock_logger.assert_called_once()
assert "An info message" in mock_logger.call_args_list[0][0][0]
@patch("logging.debug")
def test_logging_debug(self, mock_logger):
"""
Test logging_debug function.
"""
logging_debug("A debug message")
mock_logger.assert_called_once()
assert "A debug message" in mock_logger.call_args_list[0][0][0]
def test_get_context_info_normal(self):
"""
Test get_context_info function in a normal function call.
"""
def dummy_function():
return get_context_info()
module_name, line_number, function_name = dummy_function()
assert (
"test_get_context_info_normal" in function_name
) # changed from co_name to function name
assert (
"test_logging_utils.py" in module_name
) # changed from None to test_unit.py
assert isinstance(line_number, int)
def test_get_context_info_no_frame(self):
"""
Test get_context_info function when no frame is available.
"""
# Simulate a situation where inspect.currentframe() returns None (which is hard to do directly)
with patch("inspect.currentframe") as mock_frame:
mock_frame.return_value = None
module_name, line_number, function_name = get_context_info()
assert module_name is None
assert line_number is None
assert function_name is None
def test_get_context_info_exception(self):
"""
Test get_context_info function when an exception occurs.
"""
with patch("inspect.currentframe") as mock_frame:
mock_frame.side_effect = Exception("Simulated exception")
module_name, line_number, function_name = get_context_info()
assert module_name is None
assert line_number is None
assert function_name is None
tests/adapter/test_pydantic_serializer.py
+111 -17
View File
@@ -1,38 +1,132 @@
from datetime import datetime
from unittest import TestCase
from uuid import UUID
from amqp.adapter.pydantic_serializer import parse_url_query
from pydantic import BaseModel
from amqp.adapter.pydantic_serializer import (
deserialize_object,
parse_url_query,
serialize_object,
)
class NestedModel(BaseModel):
"""
Nested model for testing the pydantic serializer.
"""
value: str
number: int
class TestModel(BaseModel):
"""
Test model for testing the pydantic serializer.
"""
name: str
age: int
uuid_field: UUID
datetime_field: datetime
nested_model: NestedModel | None = None
list_field: list = []
class PydanticSerializerTest(TestCase):
def setUp(self):
"""
Prepare the test data.
"""
self.test_uuid = UUID("12345678-1234-5678-1234-567812345678")
self.test_datetime = datetime(2024, 1, 1, 12, 0, 0)
self.nested_model = NestedModel(value="test", number=42)
self.test_model = TestModel(
name="test",
age=25,
uuid_field=self.test_uuid,
datetime_field=self.test_datetime,
nested_model=self.nested_model,
list_field=[1, 2, 3],
)
def test_serialize_object(self):
assert True
serialized = serialize_object(self.test_model)
self.assertIsInstance(serialized, str)
self.assertIn("TestModel", serialized)
self.assertIn("test", serialized)
self.assertIn("25", serialized)
self.assertIn(str(self.test_uuid), serialized)
self.assertIn(self.test_datetime.isoformat(), serialized)
def test_deserialize_object(self):
assert True
serialized = serialize_object(self.test_model)
deserialized = deserialize_object(serialized)
def test__convert_value(self):
assert True
self.assertIsInstance(deserialized, TestModel)
self.assertEqual(deserialized.name, "test")
self.assertEqual(deserialized.age, 25)
self.assertEqual(deserialized.uuid_field, self.test_uuid)
self.assertEqual(deserialized.datetime_field, self.test_datetime)
self.assertIsInstance(deserialized.nested_model, NestedModel)
self.assertEqual(deserialized.nested_model.value, "test")
self.assertEqual(deserialized.nested_model.number, 42)
self.assertEqual(deserialized.list_field, [1, 2, 3])
def test__is_uuid_string(self):
assert True
def test_serialize_deserialize_with_none(self):
model = TestModel(
name="test",
age=25,
uuid_field=self.test_uuid,
datetime_field=self.test_datetime,
nested_model=None,
list_field=[],
)
serialized = serialize_object(model)
deserialized = deserialize_object(serialized)
self.assertIsNone(deserialized.nested_model)
def test__is_datetime_string(self):
assert True
def test_serialize_deserialize_with_list_of_models(self):
nested_models = [NestedModel(value=f"test{i}", number=i) for i in range(3)]
model = TestModel(
name="test",
age=25,
uuid_field=self.test_uuid,
datetime_field=self.test_datetime,
nested_model=None,
list_field=nested_models,
)
serialized = serialize_object(model)
deserialized = deserialize_object(serialized)
self.assertEqual(len(deserialized.list_field), 3)
for i, item in enumerate(deserialized.list_field):
self.assertIsInstance(item, NestedModel)
self.assertEqual(item.value, f"test{i}")
self.assertEqual(item.number, i)
def test__is_int_string(self):
assert True
def test_parse_url_query_no_params(self):
test_url = "http://example.com/path"
path, params = parse_url_query(test_url)
def test_serialize_to_json(self):
assert True
self.assertEqual(path, "/path")
self.assertEqual(params, {})
def test_parse_url_query_single_value(self):
test_url = "http://example.com/path?foo=1"
path, params = parse_url_query(test_url)
self.assertEqual(path, "/path")
self.assertEqual(params["foo"], "1")
def test_parse_url_query(self):
path, args = parse_url_query(
"http://localhost:8080/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741?response_type=JsonFile"
)
assert path == "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741"
assert args == {"response_type": "JsonFile"}
self.assertEqual(path, "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741")
self.assertEqual(args, {"response_type": "JsonFile"})
path, args = parse_url_query(
"http://localhost:8080/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741"
)
assert path == "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741"
assert args == {}
self.assertEqual(path, "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741")
self.assertEqual(args, {})
tests/adapter/test_serializer.py
+114
View File
@@ -0,0 +1,114 @@
import unittest
from io import BytesIO
from amqp.adapter.serializer import (
bytes_to_long,
csv_as_list,
long_to_bytes,
map_as_string,
map_with_list_as_string,
object_or_list_as_string,
parse_map_list_string,
parse_map_string,
read_long,
str_to_bool,
)
class TestSerializer(unittest.TestCase):
def test_long_to_bytes_and_bytes_to_long(self):
# Test conversion in both directions
test_numbers = [
-9223372036854775808,
-1,
0,
1,
9223372036854775807,
] # Min and max long values
for num in test_numbers:
bytes_val = long_to_bytes(num)
self.assertEqual(bytes_to_long(bytes_val), num)
def test_read_long(self):
# Test reading valid long
test_num = 12345
bio = BytesIO(long_to_bytes(test_num))
self.assertEqual(read_long(bio), test_num)
# Test reading from empty BytesIO
empty_bio = BytesIO()
self.assertEqual(read_long(empty_bio), 0)
def test_object_or_list_as_string(self):
# Test with string
self.assertEqual(object_or_list_as_string("test"), "test")
# Test with list
self.assertEqual(object_or_list_as_string(["a", "b", "c"]), "[a,b,c]")
# Test with empty list
self.assertEqual(object_or_list_as_string([]), "[]")
def test_map_as_string(self):
# Test with simple key-value pairs
test_map = {"key1": "value1", "key2": "value2"}
self.assertEqual(map_as_string(test_map), "key1=value1,key2=value2")
# Test with empty map
self.assertEqual(map_as_string({}), "")
def test_map_with_list_as_string(self):
# Test with lists as values
test_map = {"key1": ["a", "b"], "key2": ["c", "d"]}
self.assertEqual(map_with_list_as_string(test_map), "key1=[a,b],key2=[c,d]")
# Test with empty map
self.assertEqual(map_with_list_as_string({}), "")
def test_parse_map_string(self):
# Test valid map string
test_str = "{key1=value1,key2=value2}"
expected = {"key1": "value1", "key2": "value2"}
self.assertEqual(parse_map_string(test_str), expected)
# Test empty map string
self.assertEqual(parse_map_string("{}"), {})
# Test invalid format
self.assertEqual(parse_map_string("invalid"), {})
def test_csv_as_list(self):
# Test normal CSV string
self.assertEqual(csv_as_list("a,b,c"), ["a", "b", "c"])
# Test with spaces
self.assertEqual(csv_as_list(" a , b , c "), ["a", "b", "c"])
# Test with brackets
self.assertEqual(csv_as_list("[a,b,c]"), ["a", "b", "c"])
# Test empty list
self.assertEqual(csv_as_list("[]"), [""])
def test_parse_map_list_string(self):
# Test valid map list string
test_str = "{key1=[a,b],key2=[c,d]}"
expected = {"key1": ["a", "b"], "key2": ["c", "d"]}
self.assertEqual(parse_map_list_string(test_str), expected)
# Test empty map
self.assertEqual(parse_map_list_string("{}"), {})
# Test single entry
self.assertEqual(parse_map_list_string("{key1=[a,b]}"), {"key1": ["a", "b"]})
def test_str_to_bool(self):
# Test true values
true_values = ["true", "yes", "1", "t", "y", "on", "TRUE", "YES", "T", "Y", "ON"]
for value in true_values:
self.assertTrue(str_to_bool(value))
# Test false values
false_values = ["false", "no", "0", "f", "n", "off", "FALSE", "NO", "F", "N", "OFF"]
for value in false_values:
self.assertFalse(str_to_bool(value))
tests/adapter/test_trace_info_adapter.py
+251
View File
@@ -0,0 +1,251 @@
"""Unit tests for the tracing module."""
import unittest
from unittest.mock import MagicMock, patch
from opentelemetry.sdk.resources import Resource
from opentelemetry.trace import Tracer
from amqp.service.tracing import (
MetricsConfig,
TracingConfig,
create_resource,
create_span,
initialize_telemetry,
setup_metrics_provider,
setup_otlp_exporter,
setup_trace_provider,
)
class TestTracingConfig(unittest.TestCase):
"""Test cases for TracingConfig."""
def test_default_values(self):
"""Test that default values are set correctly."""
config = TracingConfig(service_name="test-service", service_version="1.0.0")
self.assertEqual(config.service_name, "test-service")
self.assertEqual(config.service_version, "1.0.0")
self.assertEqual(config.otlp_endpoint, "http://localhost:4317")
self.assertTrue(config.insecure)
self.assertFalse(config.debug)
self.assertIsNone(config.additional_attributes)
def test_custom_values(self):
"""Test that custom values are set correctly."""
config = TracingConfig(
service_name="test-service",
service_version="1.0.0",
otlp_endpoint="http://custom:4317",
insecure=False,
debug=True,
additional_attributes={"custom": "attribute"},
)
self.assertEqual(config.otlp_endpoint, "http://custom:4317")
self.assertFalse(config.insecure)
self.assertTrue(config.debug)
self.assertEqual(config.additional_attributes, {"custom": "attribute"})
class TestMetricsConfig(unittest.TestCase):
"""Test cases for MetricsConfig."""
def test_default_values(self):
"""Test that default values are set correctly."""
config = MetricsConfig()
self.assertEqual(config.prometheus_port, 9464)
self.assertEqual(config.prometheus_host, "localhost")
self.assertEqual(config.export_interval_millis, 30000)
self.assertIsNone(config.additional_attributes)
def test_custom_values(self):
"""Test that custom values are set correctly."""
config = MetricsConfig(
prometheus_port=8000,
prometheus_host="0.0.0.0",
export_interval_millis=15000,
additional_attributes={"custom": "attribute"},
)
self.assertEqual(config.prometheus_port, 8000)
self.assertEqual(config.prometheus_host, "0.0.0.0")
self.assertEqual(config.export_interval_millis, 15000)
self.assertEqual(config.additional_attributes, {"custom": "attribute"})
class TestCreateResource(unittest.TestCase):
"""Test cases for create_resource function."""
def test_create_resource_basic(self):
"""Test creating a resource with basic configuration."""
config = TracingConfig(service_name="test-service", service_version="1.0.0")
resource = create_resource(config)
self.assertIsInstance(resource, Resource)
self.assertEqual(resource.attributes.get("service.name"), "test-service")
self.assertEqual(resource.attributes.get("service.version"), "1.0.0")
def test_create_resource_with_attributes(self):
"""Test creating a resource with additional attributes."""
config = TracingConfig(
service_name="test-service",
service_version="1.0.0",
additional_attributes={"custom": "attribute"},
)
resource = create_resource(config)
self.assertEqual(resource.attributes.get("custom"), "attribute")
class TestSetupExporters(unittest.TestCase):
"""Test cases for exporter setup functions."""
@patch("amqp.service.tracing.OTLPSpanExporter")
def test_setup_otlp_exporter(self, mock_otlp):
"""Test OTLP exporter setup."""
config = TracingConfig(
service_name="test-service",
service_version="1.0.0",
otlp_endpoint="http://jaeger:4317",
insecure=True,
)
setup_otlp_exporter(config)
mock_otlp.assert_called_once_with(endpoint="http://jaeger:4317", insecure=True)
@patch("amqp.service.tracing.OTLPSpanExporter")
def test_setup_otlp_exporter_secure(self, mock_otlp):
"""Test OTLP exporter setup with secure connection."""
config = TracingConfig(
service_name="test-service",
service_version="1.0.0",
otlp_endpoint="https://jaeger:4317",
insecure=False,
)
setup_otlp_exporter(config)
mock_otlp.assert_called_once_with(endpoint="https://jaeger:4317", insecure=False)
class TestSetupTraceProvider(unittest.TestCase):
"""Test cases for trace provider setup."""
@patch("amqp.service.tracing.TracerProvider")
@patch("amqp.service.tracing.BatchSpanProcessor")
@patch("amqp.service.tracing.setup_otlp_exporter")
def test_setup_trace_provider_basic(self, mock_otlp_exporter, mock_processor, mock_provider):
"""Test basic trace provider setup."""
config = TracingConfig(service_name="test-service", service_version="1.0.0")
resource = create_resource(config)
mock_provider_instance = MagicMock()
mock_provider.return_value = mock_provider_instance
setup_trace_provider(config, resource)
mock_provider.assert_called_once_with(resource=resource)
mock_otlp_exporter.assert_called_once()
self.assertEqual(mock_provider_instance.add_span_processor.call_count, 1)
@patch("amqp.service.tracing.TracerProvider")
@patch("amqp.service.tracing.BatchSpanProcessor")
@patch("amqp.service.tracing.setup_otlp_exporter")
@patch("amqp.service.tracing.ConsoleSpanExporter")
def test_setup_trace_provider_with_debug(
self, mock_console_exporter, mock_otlp_exporter, mock_processor, mock_provider
):
"""Test trace provider setup with debug mode."""
config = TracingConfig(service_name="test-service", service_version="1.0.0", debug=True)
resource = create_resource(config)
mock_provider_instance = MagicMock()
mock_provider.return_value = mock_provider_instance
setup_trace_provider(config, resource)
self.assertEqual(mock_provider_instance.add_span_processor.call_count, 2)
mock_console_exporter.assert_called_once()
class TestSetupMetricsProvider(unittest.TestCase):
"""Test cases for metrics provider setup."""
@patch("amqp.service.tracing.start_http_server")
@patch("amqp.service.tracing.PrometheusMetricReader")
@patch("amqp.service.tracing.PeriodicExportingMetricReader")
@patch("amqp.service.tracing.MeterProvider")
def test_setup_metrics_provider_success(
self, mock_meter_provider, mock_periodic_reader, mock_prometheus_reader, mock_start_server
):
"""Test successful metrics provider setup."""
config = MetricsConfig(prometheus_port=9464, prometheus_host="localhost")
resource = Resource({})
setup_metrics_provider(config, resource)
mock_start_server.assert_called_once_with(port=9464, addr="localhost")
mock_prometheus_reader.assert_called_once()
mock_periodic_reader.assert_called_once()
mock_meter_provider.assert_called_once()
@patch("amqp.service.tracing.start_http_server")
def test_setup_metrics_provider_failure(self, mock_start_server):
"""Test metrics provider setup with failure."""
mock_start_server.side_effect = Exception("Failed to start server")
config = MetricsConfig()
resource = Resource({})
provider = setup_metrics_provider(config, resource)
self.assertIsNotNone(provider) # Should return no-op provider
class TestInitializeTelemetry(unittest.TestCase):
"""Test cases for initialize_telemetry function."""
@patch("amqp.service.tracing.setup_trace_provider")
@patch("amqp.service.tracing.setup_metrics_provider")
@patch("amqp.service.tracing.trace")
@patch("amqp.service.tracing.metrics")
def test_initialize_telemetry_success(
self, mock_metrics, mock_trace, mock_setup_metrics, mock_setup_trace
):
"""Test successful telemetry initialization."""
config = TracingConfig(service_name="test-service", service_version="1.0.0")
metrics_config = MetricsConfig()
mock_tracer = MagicMock(spec=Tracer)
mock_trace.get_tracer.return_value = mock_tracer
tracer = initialize_telemetry(config, metrics_config)
self.assertEqual(tracer, mock_tracer)
mock_setup_trace.assert_called_once()
mock_setup_metrics.assert_called_once()
@patch("amqp.service.tracing.setup_trace_provider")
def test_initialize_telemetry_failure(self, mock_setup_trace):
"""Test telemetry initialization with failure."""
mock_setup_trace.side_effect = Exception("Setup failed")
tracer = initialize_telemetry()
self.assertIsNotNone(tracer) # Should return no-op tracer
class TestCreateSpan(unittest.TestCase):
"""Test cases for create_span function."""
def test_create_span_basic(self):
"""Test basic span creation."""
mock_tracer = MagicMock(spec=Tracer)
create_span(mock_tracer, "test-span")
mock_tracer.start_as_current_span.assert_called_once_with("test-span", attributes={})
def test_create_span_with_attributes(self):
"""Test span creation with attributes."""
mock_tracer = MagicMock(spec=Tracer)
attributes = {"custom": "attribute"}
create_span(mock_tracer, "test-span", attributes)
mock_tracer.start_as_current_span.assert_called_once_with(
"test-span", attributes=attributes
)
if __name__ == "__main__":
unittest.main()
tests/model/test_model.py
+220
View File
@@ -0,0 +1,220 @@
import base64
import json
import unittest
from amqp.model.model import (
AMQErrorMessage,
AMQRoute,
CleverMicroMessage,
DataMessage,
DataResponse,
ScalingRequest,
ScalingRequestAlert,
)
from amqp.model.snowflake_id import SnowflakeId
class TestCleverMicroMessage(unittest.TestCase):
def setUp(self):
self.test_id = SnowflakeId()
self.test_headers = {"Content-Type": ["application/json"]}
self.test_trace_info = {"trace_id": "123"}
self.test_body = b"test data"
self.test_base64body = base64.b64encode(self.test_body)
self.message = CleverMicroMessage(
magic="A",
id=self.test_id,
m_field="GET",
d_field="test_domain",
p_field="test_path",
headers=self.test_headers,
trace_info=self.test_trace_info,
base64body=self.test_base64body,
)
def test_getters(self):
self.assertEqual(self.message.magic(), "A")
self.assertEqual(self.message.id(), self.test_id)
self.assertEqual(self.message.m_field(), "GET")
self.assertEqual(self.message.d_field(), "test_domain")
self.assertEqual(self.message.p_field(), "test_path")
self.assertEqual(self.message.headers(), self.test_headers)
self.assertEqual(self.message.trace_info(), self.test_trace_info)
self.assertEqual(self.message.base64body(), self.test_base64body)
def test_body_decoding(self):
self.assertEqual(self.message.body(), self.test_body)
def test_content_type(self):
self.assertEqual(self.message.content_type(), "application/json")
# Test default content type
message_no_content_type = CleverMicroMessage(
magic="A",
id=self.test_id,
m_field="GET",
d_field="test_domain",
p_field="test_path",
headers={},
trace_info={},
base64body=self.test_base64body,
)
self.assertEqual(
message_no_content_type.content_type(), CleverMicroMessage.DEFAULT_CONTENT_TYPE[0]
)
class TestDataMessage(unittest.TestCase):
def setUp(self):
self.test_id = SnowflakeId()
self.test_headers = {"Content-Type": ["application/json"]}
self.test_trace_info = {"trace_id": "123"}
self.test_base64body = base64.b64encode(b"test data")
self.message = DataMessage(
magic="A",
id=self.test_id,
method="GET",
domain="testdomain",
path="testpath",
headers=self.test_headers,
trace_info=self.test_trace_info,
base64body=self.test_base64body,
)
def test_getters(self):
self.assertEqual(self.message.method(), "GET")
self.assertEqual(self.message.domain(), "testdomain")
self.assertEqual(self.message.path(), "testpath")
def test_routing_key(self):
self.assertEqual(self.message.routing_key(), "testdomain.testpath")
class TestDataResponse(unittest.TestCase):
def setUp(self):
self.test_id = SnowflakeId()
self.test_trace_info = {"trace_id": "123"}
self.test_base64body = base64.b64encode(b"test data")
self.response = DataResponse(
magic="R",
id=self.test_id,
response_code="200",
content_type="application/json",
error="",
error_cause="",
headers={},
trace_info=self.test_trace_info,
base64body=self.test_base64body,
)
def test_getters(self):
self.assertEqual(self.response.response_code(), "200")
self.assertEqual(self.response.error(), "")
self.assertEqual(self.response.error_cause(), "")
self.assertEqual(self.response.content_type(), "application/json")
class TestAMQRoute(unittest.TestCase):
def setUp(self):
self.route = AMQRoute(
component_name="test_component",
exchange="test_exchange",
queue="test_queue",
key="test_key",
timeout=1000,
valid_until=1735689600000, # Some future timestamp
)
def test_str_representation(self):
expected = "test_component:test_exchange:test_queue:test_key:1000:1735689600000"
self.assertEqual(str(self.route), expected)
def test_equality(self):
same_route = AMQRoute(
component_name="test_component",
exchange="test_exchange",
queue="test_queue",
key="test_key",
timeout=1000,
valid_until=1735689600000,
)
different_route = AMQRoute(
component_name="other_component",
exchange="test_exchange",
queue="test_queue",
key="test_key",
timeout=1000,
valid_until=1735689600000,
)
self.assertEqual(self.route, same_route)
self.assertNotEqual(self.route, different_route)
class TestScalingRequest(unittest.TestCase):
def setUp(self):
self.request = ScalingRequest(
service_id="test_service",
task_id="test_task",
max_availability=5,
current_availability=3,
request_type=ScalingRequestAlert.OVERLOAD,
version=1,
)
def test_to_json(self):
json_str = self.request.to_json()
data = json.loads(json_str)
self.assertEqual(data["serviceId"], "test_service")
self.assertEqual(data["taskId"], "test_task")
self.assertEqual(data["maxAvailability"], 5)
self.assertEqual(data["currentAvailability"], 3)
self.assertEqual(data["requestType"], "OVERLOAD")
self.assertEqual(data["version"], 1)
def test_from_json(self):
json_str = json.dumps(
{
"serviceId": "test_service",
"taskId": "test_task",
"maxAvailability": 5,
"currentAvailability": 3,
"requestType": "OVERLOAD",
"version": 1,
}
)
request = ScalingRequest.from_json(json_str)
self.assertEqual(request.service_id, "test_service")
self.assertEqual(request.task_id, "test_task")
self.assertEqual(request.max_availability, 5)
self.assertEqual(request.current_availability, 3)
self.assertEqual(request.request_type, ScalingRequestAlert.OVERLOAD)
self.assertEqual(request.version, 1)
class TestAMQErrorMessage(unittest.TestCase):
def setUp(self):
self.error = AMQErrorMessage(
response_code=400,
error="Bad Request",
detail="Invalid input",
content_type="application/json",
)
def test_serialize(self):
serialized = self.error.serialize()
data = json.loads(serialized)
self.assertEqual(data["response_code"], 400)
self.assertEqual(data["error"], "Bad Request")
self.assertEqual(data["detail"], "Invalid input")
if __name__ == "__main__":
unittest.main()
tests/rabbitmq/test_rabbit_mq_client.py
+209
View File
@@ -0,0 +1,209 @@
from unittest import IsolatedAsyncioTestCase
from unittest.mock import AsyncMock, MagicMock, patch
from amqp.model.model import DataMessage
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient
class TestRabbitMQClient(IsolatedAsyncioTestCase):
def setUp(self):
self.mock_config = MagicMock()
self.mock_config.dispatch.amq_host = "localhost"
self.mock_config.dispatch.amq_port = 5672
self.mock_config.dispatch.rabbit_mq_user = "user"
self.mock_config.dispatch.rabbit_mq_password = "pass"
self.mock_config.dispatch.use_dlq = False
patcher = patch("amqp.rabbitmq.rabbit_mq_client.RouterConsumer")
self.addCleanup(patcher.stop)
self.mock_router = patcher.start()
self.client = RabbitMQClient(self.mock_config)
async def test_init_sets_up_channel_and_reply_exchange(self):
"""Test that init sets up the channel, reply queue, and reply exchange."""
loop = MagicMock()
self.client.setup_amq_channel = AsyncMock()
self.client.router.init_routing_paths_consumer = AsyncMock()
self.client.router.get_reply_to_queue_name = MagicMock(return_value="reply_queue")
self.client.channel = MagicMock()
self.client.channel.declare_queue = AsyncMock()
self.client.channel.declare_exchange = AsyncMock()
queue_mock = MagicMock()
exchange_mock = MagicMock()
self.client.channel.declare_queue.return_value = queue_mock
self.client.channel.declare_exchange.return_value = exchange_mock
queue_mock.bind = AsyncMock()
result = await self.client.init(loop)
self.client.setup_amq_channel.assert_awaited_once_with(loop)
self.client.router.init_routing_paths_consumer.assert_awaited_once()
self.client.channel.declare_queue.assert_awaited_once()
self.client.channel.declare_exchange.assert_awaited_once()
queue_mock.bind.assert_awaited_once_with(exchange=exchange_mock, routing_key="reply_queue")
self.assertEqual(result, exchange_mock)
def test_cleanup_closes_channel_and_connection(self):
"""Test that cleanup closes the channel and connection if open."""
self.client.router.cleanup = MagicMock()
self.client.router.is_open = MagicMock(return_value=True)
self.client.channel = MagicMock()
self.client.connection = MagicMock()
self.client.connection.is_closed = False
self.client.channel.close = MagicMock()
self.client.connection.close = MagicMock()
self.client.cleanup()
self.client.router.cleanup.assert_called_once()
self.client.channel.close.assert_called_once()
self.client.connection.close.assert_called_once()
def test_cleanup_handles_exceptions(self):
"""Test that cleanup handles exceptions during close."""
self.client.router.cleanup = MagicMock()
self.client.router.is_open = MagicMock(return_value=True)
self.client.channel = MagicMock()
self.client.connection = MagicMock()
self.client.connection.is_closed = False
self.client.channel.close = MagicMock(side_effect=Exception("fail"))
self.client.connection.close = MagicMock()
# Should not raise
self.client.cleanup()
async def test_register_inbound_routes_registers_routes(self):
"""Test that register_inbound_routes declares queues, exchanges, binds, and registers routes."""
self.client.router.is_open = MagicMock(return_value=True)
self.client.amq_configuration = self.mock_config
self.client.router.unwrap_route_list = MagicMock(
return_value=[
MagicMock(
queue="q1",
exchange="ex1",
key="rk1",
component_name="comp",
timeout=0,
valid_until=0,
)
]
)
self.client.router.get_consume_queue_name = MagicMock(return_value="q1")
self.client.channel = MagicMock()
self.client.channel.declare_queue = AsyncMock()
self.client.channel.declare_exchange = AsyncMock()
queue_mock = MagicMock()
exchange_mock = MagicMock()
self.client.channel.declare_queue.return_value = queue_mock
self.client.channel.declare_exchange.return_value = exchange_mock
queue_mock.bind = AsyncMock()
queue_mock.consume = AsyncMock(return_value="ctag")
self.client.router.register_route = AsyncMock()
await self.client.register_inbound_routes("route_map", MagicMock())
self.client.channel.declare_queue.assert_awaited_once()
self.client.channel.declare_exchange.assert_awaited_once()
queue_mock.bind.assert_awaited_once_with(exchange_mock, "rk1")
queue_mock.consume.assert_awaited_once()
self.client.router.register_route.assert_awaited_once()
async def test_register_inbound_routes_channel_null(self):
"""Test that register_inbound_routes logs a warning if channel is null."""
self.client.router.is_open = MagicMock(return_value=False)
with patch("amqp.rabbitmq.rabbit_mq_client.logging_warning") as mock_log:
await self.client.register_inbound_routes("route_map", MagicMock())
mock_log.assert_called()
async def test_register_data_reply_callback(self):
"""Test that register_data_reply_callback consumes on reply_to_queue."""
queue_mock = MagicMock()
queue_mock.consume = AsyncMock(return_value="ctag")
queue_mock.name = "reply_queue"
self.client.reply_to_queue = queue_mock
with patch("amqp.rabbitmq.rabbit_mq_client.logging_info") as mock_log:
await self.client.register_data_reply_callback(MagicMock())
queue_mock.consume.assert_awaited_once()
mock_log.assert_called()
async def test_setup_amq_channel_sets_channel_and_qos(self):
"""Test that setup_amq_channel sets up the connection, channel, and QoS."""
loop = MagicMock()
with patch("aio_pika.connect_robust", new_callable=AsyncMock) as mock_connect:
mock_connection = MagicMock()
mock_connect.return_value = mock_connection
mock_connection.channel = AsyncMock()
mock_channel = MagicMock()
mock_connection.channel.return_value = mock_channel
mock_channel.set_qos = AsyncMock()
self.client.router.init_internal_routing_paths = AsyncMock()
self.client.router.set_route_added_notifier = MagicMock()
await self.client.setup_amq_channel(loop)
mock_connect.assert_awaited_once()
mock_connection.channel.assert_awaited_once()
mock_channel.set_qos.assert_awaited_once_with(
prefetch_count=RabbitMQClient.PREFETCH_COUNT
)
self.client.router.init_internal_routing_paths.assert_awaited_once_with(mock_channel)
self.client.router.set_route_added_notifier.assert_called_once()
async def test_send_message_no_route(self):
"""Test that send_message logs a warning if no route is found."""
self.client.router.find_route_by_message = MagicMock(return_value=None)
self.client.rpc_exchange = MagicMock()
message = MagicMock(spec=DataMessage)
with patch("amqp.rabbitmq.rabbit_mq_client.logging_warning") as mock_log:
await self.client.send_message(message)
mock_log.assert_called()
async def test_send_message_no_reply(self):
"""Test that send_message publishes a message with no reply expected."""
route = MagicMock(timeout=0)
self.client.router.find_route_by_message = MagicMock(return_value=route)
self.client.rpc_exchange = AsyncMock()
self.client.router.get_routing_key = MagicMock(return_value="rk")
message = MagicMock(spec=DataMessage)
with patch(
"amqp.rabbitmq.rabbit_mq_client.DataMessageFactory.serialize", return_value=b"bytes"
):
with patch("amqp.rabbitmq.rabbit_mq_client.logging_info") as mock_log:
await self.client.send_message(message)
self.client.rpc_exchange.publish.assert_awaited_once()
mock_log.assert_called()
async def test_send_message_rpc_call(self):
"""Test that send_message publishes a message with reply expected (RPC)."""
route = MagicMock(timeout=1)
self.client.router.find_route_by_message = MagicMock(return_value=route)
self.client.rpc_exchange = AsyncMock()
self.client.router.get_routing_key = MagicMock(return_value="rk")
self.client.router.get_reply_to_queue_name = MagicMock(return_value="reply_q")
message = MagicMock(spec=DataMessage)
message.id.return_value = "id"
with patch(
"amqp.rabbitmq.rabbit_mq_client.DataMessageFactory.serialize", return_value=b"bytes"
):
await self.client.send_message(message)
self.client.rpc_exchange.publish.assert_awaited_once()
async def test_send_message_exception(self):
"""Test that send_message logs an error if an exception occurs."""
self.client.router.find_route_by_message = MagicMock(side_effect=Exception("fail"))
message = MagicMock(spec=DataMessage)
with patch("amqp.rabbitmq.rabbit_mq_client.logging_error") as mock_log:
await self.client.send_message(message)
mock_log.assert_called()
async def test_request_routes_delegates(self):
"""Test that request_routes delegates to router."""
self.client.router.request_routes_from_adapters = AsyncMock()
await self.client.request_routes()
self.client.router.request_routes_from_adapters.assert_awaited_once()
def test_find_route_delegates(self):
"""Test that find_route delegates to router."""
self.client.router.find_route = MagicMock(return_value="route")
result = self.client.find_route("domain", "path")
self.client.router.find_route.assert_called_once_with("domain", "path")
self.assertEqual(result, "route")
tests/requirements-test.txt
+4
View File
@@ -0,0 +1,4 @@
pytest>=7.0.0
pytest-cov>=4.0.0
pytest-asyncio>=0.21.0
aio-pika>=9.0.0
tests/router/test_router_consumer.py
+210
View File
@@ -0,0 +1,210 @@
import unittest
from unittest.mock import AsyncMock, MagicMock, patch
from aio_pika.abc import (
AbstractIncomingMessage,
AbstractRobustChannel,
AbstractRobustQueue,
)
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQRoute, ServiceMessage
from amqp.model.service_message_type import ServiceMessageType
from amqp.router.router_base import (
CM_C_REQ_QUEUE,
CM_REQUEST_EXCHANGE,
)
from amqp.router.router_consumer import RouterConsumer
class TestRouterConsumer(unittest.IsolatedAsyncioTestCase):
def setUp(self):
# Create a mock for amq_adapter
self.amq_adapter = MagicMock()
self.amq_adapter.adapter_prefix.return_value = "test_"
self.amq_adapter.service_name = "test_service"
# Create the config mock with amq_adapter
self.config = MagicMock(spec=AMQConfiguration)
self.config.amq_adapter = self.amq_adapter
self.consumer = RouterConsumer(self.config)
# Mock channel
self.channel = AsyncMock(spec=AbstractRobustChannel)
self.consumer.channel = self.channel
# Mock queue
self.queue = AsyncMock(spec=AbstractRobustQueue)
self.channel.declare_queue.return_value = self.queue
def test_init(self):
"""Test initialization of RouterConsumer"""
self.assertIsNotNone(self.consumer)
self.assertEqual(self.consumer.amq_configuration, self.config)
@patch("amqp.router.router_consumer.logging_info")
async def test_init_routing_paths_consumer_success(self, mock_logging):
"""Test successful initialization of routing paths consumer"""
self.consumer.is_open = MagicMock(return_value=True)
await self.consumer.init_routing_paths_consumer(self.channel)
# Verify queue declaration
self.channel.declare_queue.assert_called_once_with(
name="test_" + CM_C_REQ_QUEUE, durable=True, exclusive=False, auto_delete=False
)
# Verify queue binding
self.queue.bind.assert_called_once_with(exchange=CM_REQUEST_EXCHANGE, routing_key="")
# Verify consumer setup
self.queue.consume.assert_called_once()
@patch("amqp.router.router_consumer.logging_error")
async def test_init_routing_paths_consumer_channel_closed(self, mock_logging):
"""Test initialization when channel is closed"""
self.consumer.is_open = MagicMock(return_value=False)
await self.consumer.init_routing_paths_consumer(self.channel)
mock_logging.assert_called_once()
self.channel.declare_queue.assert_not_called()
@patch("amqp.router.router_consumer.logging_debug")
@patch("amqp.router.router_consumer.logging_info")
async def test_handle_routing_data_request_message_valid(
self, mock_logging_info, mock_logging_debug
):
"""Test handling of valid routing data request message"""
# Setup
message = AsyncMock(spec=AbstractIncomingMessage)
message.delivery_tag = 1
message.consumer_tag = "test_consumer"
message.properties = MagicMock()
message.body = b"test_body"
# Mock service message factory
mock_service_message = ServiceMessage(
id=None,
payload_type=0,
message_type=ServiceMessageType.ROUTING_DATA_REQUEST,
recipient_name="test_service",
reply_to="other_service",
trace_info={},
message="test_payload",
)
self.consumer.service_message_factory = MagicMock()
self.consumer.service_message_factory.from_bytes = MagicMock(
return_value=mock_service_message
)
self.consumer.service_message_factory.of = MagicMock()
self.consumer.wrap_own_routes = MagicMock(return_value="test_routes")
self.consumer.publish_service_message = AsyncMock(return_value=True)
# Execute
await self.consumer.handle_routing_data_request_message(message)
# Verify
message.ack.assert_called_once_with(False)
self.consumer.publish_service_message.assert_called_once()
@patch("amqp.router.router_consumer.logging_info")
async def test_register_route_success(self, mock_logging):
"""Test successful route registration"""
# Setup
route = AMQRoute(
component_name="test_service",
exchange="test_exchange",
queue="test_queue",
key="test_key",
timeout=1000,
valid_until=0,
)
self.consumer.is_open = MagicMock(return_value=True)
self.consumer.publish_service_message = AsyncMock(return_value=True)
self.consumer.add_own_route = MagicMock()
# Execute
await self.consumer.register_route(route)
# Verify
self.consumer.publish_service_message.assert_called_once()
self.consumer.add_own_route.assert_called_once_with(route)
@patch("amqp.router.router_consumer.logging_warning")
async def test_register_route_channel_closed(self, mock_logging):
"""Test route registration when channel is closed"""
# Setup
route = AMQRoute(
component_name="test_service",
exchange="test_exchange",
queue="test_queue",
key="test_key",
timeout=1000,
valid_until=0,
)
self.consumer.is_open = MagicMock(return_value=False)
# Execute
await self.consumer.register_route(route)
# Verify
mock_logging.assert_called_once()
@patch("amqp.router.router_consumer.logging_info")
@patch("amqp.router.router_consumer.logging_warning")
async def test_remove_route_on_cleanup(self, mock_logging_warning, mock_logging_info):
"""Test route removal during cleanup"""
# Setup
route = AMQRoute(
component_name="test_service",
exchange="test_exchange",
queue="test_queue",
key="test_key",
timeout=1000,
valid_until=0,
)
self.consumer.publish_service_message = AsyncMock(return_value=True)
self.consumer.remove_route = MagicMock(return_value=True)
# Execute
result = await self.consumer.remove_route_on_cleanup(route)
# Verify
self.consumer.publish_service_message.assert_called_once()
self.consumer.remove_route.assert_called_once_with(route)
self.assertTrue(result)
@patch("amqp.router.router_consumer.logging_info")
@patch("amqp.router.router_consumer.logging_warning")
def test_cleanup(self, mock_logging_warning, mock_logging_info):
"""Test cleanup process"""
# Setup
routes = {
AMQRoute(
component_name="test_service",
exchange="test_exchange1",
queue="test_queue1",
key="test_key1",
timeout=1000,
valid_until=0,
),
AMQRoute(
component_name="test_service",
exchange="test_exchange2",
queue="test_queue2",
key="test_key2",
timeout=1000,
valid_until=0,
),
}
self.consumer.get_routes = MagicMock(return_value=routes)
self.consumer.remove_route_on_cleanup = AsyncMock(return_value=True)
self.consumer.wrap_routes = MagicMock(return_value="test_routes")
# Execute
self.consumer.cleanup()
# Verify
self.assertEqual(self.consumer.remove_route_on_cleanup.call_count, 2)
mock_logging_info.assert_called()
tests/router/test_router_producer.py
+164
View File
@@ -0,0 +1,164 @@
import unittest
from unittest.mock import AsyncMock, MagicMock, patch
from amqp.model.service_message_type import ServiceMessageType
from amqp.router.router_producer import RouterProducer
class DummyAMQAdapter:
def __init__(self):
self.service_name = "test-service"
self.route_mapping = "route-map"
def adapter_prefix(self):
return "prefix-"
class DummyConfig:
def __init__(self):
self.amq_adapter = DummyAMQAdapter()
class TestRouterProducer(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.config = DummyConfig()
self.router = RouterProducer(self.config)
def test_get_reply_to_queue_name(self):
self.assertEqual(self.router.get_reply_to_queue_name(), "prefix-P-REPLY-TO")
def test_get_response_queue_name(self):
self.assertEqual(self.router.get_response_queue_name(), "prefix-P-RESP")
def test_get_consume_queue_name(self):
self.assertEqual(self.router.get_consume_queue_name(), "prefix-C-AMQ")
def test_is_open_true(self):
channel = MagicMock()
channel.is_closed = False
self.assertTrue(self.router.is_open(channel))
def test_is_open_false(self):
channel = MagicMock()
channel.is_closed = True
self.assertFalse(self.router.is_open(channel))
self.assertFalse(self.router.is_open(None))
@patch("amqp.router.router_producer.logging_info")
@patch("amqp.router.router_producer.logging_error")
async def test_init_internal_routing_paths_channel_open(self, mock_error, mock_info):
channel = AsyncMock()
channel.is_closed = False
channel.declare_exchange = AsyncMock()
channel.declare_queue = AsyncMock()
queue = AsyncMock()
channel.declare_queue.return_value = queue
queue.bind = AsyncMock()
queue.consume = AsyncMock(return_value="ctag")
await self.router.init_internal_routing_paths(channel)
self.assertIsNotNone(self.router.cm_request_exchange)
self.assertIsNotNone(self.router.cm_response_exchange)
queue.bind.assert_awaited()
queue.consume.assert_awaited()
mock_error.assert_not_called()
@patch("amqp.router.router_producer.logging_error")
@patch("amqp.router.router_producer.logging_info")
async def test_init_internal_routing_paths_channel_closed(self, mock_info, mock_error):
channel = MagicMock()
channel.is_closed = True
await self.router.init_internal_routing_paths(channel)
mock_error.assert_called()
@patch("amqp.router.router_producer.logging_info")
@patch("amqp.router.router_producer.logging_debug")
async def test_handle_routing_data_message_valid(self, mock_debug, mock_info):
message = AsyncMock()
message.delivery_tag = 1
message.body = b"body"
message.message_id = "mid"
message.ack = AsyncMock()
# Patch service_message_factory
service_message = MagicMock()
service_message.payload_type = 0
service_message.is_valid.return_value = True
service_message.message_type = ServiceMessageType.ROUTING_DATA
service_message.reply_to = "other-service"
service_message.message = "route-string"
self.router.service_message_factory.from_bytes = MagicMock(return_value=service_message)
self.router.amq_configuration.amq_adapter.service_name = "test-service"
self.router.add_routes = MagicMock()
self.router.remove_routes = MagicMock()
await self.router.handle_routing_data_message(message)
message.ack.assert_awaited()
self.router.add_routes.assert_called_with("route-string")
self.router.remove_routes.assert_not_called()
@patch("amqp.router.router_producer.logging_info")
@patch("amqp.router.router_producer.logging_debug")
async def test_handle_routing_data_message_remove(self, mock_debug, mock_info):
message = AsyncMock()
message.delivery_tag = 1
message.body = b"body"
message.message_id = "mid"
message.ack = AsyncMock()
service_message = MagicMock()
service_message.is_valid.return_value = True
service_message.message_type = ServiceMessageType.ROUTING_DATA_REMOVE
service_message.payload_type = 0
service_message.reply_to = "other-service"
service_message.message = "route-string"
self.router.service_message_factory.from_bytes = MagicMock(return_value=service_message)
self.router.amq_configuration.amq_adapter.service_name = "test-service"
self.router.add_routes = MagicMock()
self.router.remove_routes = MagicMock()
await self.router.handle_routing_data_message(message)
message.ack.assert_awaited()
self.router.remove_routes.assert_called_with("route-string")
self.router.add_routes.assert_not_called()
@patch("amqp.router.router_producer.logging_info")
@patch("amqp.router.router_producer.logging_warning")
@patch("amqp.router.router_producer.logging_error")
async def test_request_routes_from_adapters_success(self, mock_error, mock_warning, mock_info):
self.router.cm_request_exchange = AsyncMock()
sm = MagicMock()
self.router.service_message_factory.of = MagicMock(return_value=sm)
self.router.publish_service_message = AsyncMock(return_value=True)
await self.router.request_routes_from_adapters()
mock_warning.assert_not_called()
mock_error.assert_not_called()
@patch("amqp.router.router_producer.logging_info")
@patch("amqp.router.router_producer.logging_warning")
@patch("amqp.router.router_producer.logging_error")
async def test_request_routes_from_adapters_fail(self, mock_error, mock_warning, mock_info):
self.router.cm_request_exchange = AsyncMock()
sm = MagicMock()
self.router.service_message_factory.of = MagicMock(return_value=sm)
self.router.publish_service_message = AsyncMock(return_value=False)
await self.router.request_routes_from_adapters()
mock_warning.assert_called()
mock_error.assert_not_called()
@patch("amqp.router.router_producer.logging_info")
async def test_publish_service_message_open(self, mock_info):
self.router.channel = MagicMock()
self.router.channel.is_closed = False
self.router.service_message_factory.serialize = MagicMock(return_value=b"bin")
exchange = AsyncMock()
exchange.publish = AsyncMock()
message = MagicMock()
with patch("amqp.router.router_producer.Message") as MockMessage:
MockMessage.return_value = MagicMock()
result = await self.router.publish_service_message(message, exchange)
exchange.publish.assert_awaited()
self.assertTrue(result)
async def test_publish_service_message_closed(self):
self.router.channel = MagicMock()
self.router.channel.is_closed = True
exchange = AsyncMock()
message = MagicMock()
result = await self.router.publish_service_message(message, exchange)
self.assertFalse(result)
tests/service/test_amq_service.py
+185
View File
@@ -0,0 +1,185 @@
import unittest
from unittest.mock import AsyncMock, MagicMock, Mock, patch
from amqp.config.amq_configuration import AMQAdapter, AMQConfiguration
from amqp.model.model import DataMessage
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.generator_id = 12345
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_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_threads = 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_send_message(self):
# Create mock message and future
mock_message = Mock(spec=DataMessage)
mock_message.id.return_value = "test-id"
mock_future = MagicMock()
# Setup mock message handler
self.service.amq_data_message_handler = Mock()
self.service.amq_data_message_handler.outstanding = {}
# Send message
self.service.send_message(mock_message, mock_future)
# Verify message was added to outstanding
self.assertEqual(self.service.amq_data_message_handler.outstanding["test-id"], mock_future)
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()
tests/test_amq_configuration.py
+12
View File
@@ -0,0 +1,12 @@
class AMQConfigurationTest:
def test_amqadapter(self):
assert False
def test_dispatch(self):
assert False
def test_backpressure(self):
assert False
def test_amqconfiguration(self):
assert False