Compare commits

..

1 Commits

Author SHA1 Message Date
hurui200320 248469fe1e fix(General): Fix deps version so cleverswarm can install it
Unit test coverage / pytest (push) Failing after 1m0s
2025-04-30 21:37:26 +08:00
66 changed files with 1277 additions and 6860 deletions
-118
View File
@@ -1,118 +0,0 @@
---
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 -=
-4
View File
@@ -1,4 +0,0 @@
[flake8]
max-line-length = 88
ignore = E501, W503, E203
exclude = .git,__pycache__,docs/source/conf.py,old,build,dist,.venv
+3 -3
View File
@@ -40,9 +40,9 @@ jobs:
# publish
- name: publish to pulp pypi
run: |
twine upload --repository-url https://git.cleverthis.com/api/packages/clevermicro/pypi \
--username ${CI_REGISTRY_USER} \
--password ${CI_REGISTRY_PASSWORD} \
twine upload --repository-url https://pkg.cleverthis.io/pypi/clevermicro/amq-adapter-python/simple/ \
--username $CI_REGISTRY_USER \
--password $CI_REGISTRY_PASSWORD \
--verbose \
dist/*
env:
-21
View File
@@ -1,21 +0,0 @@
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.2.0
hooks:
- id: check-yaml
- id: end-of-file-fixer
- id: trailing-whitespace
- repo: https://github.com/PyCQA/isort
rev: 6.0.1
hooks:
- id: isort
args: ["--profile", "black"]
- repo: https://github.com/psf/black
rev: 25.1.0
hooks:
- id: black
language_version: python3.11
- repo: https://github.com/PyCQA/flake8
rev: 7.2.0
hooks:
- id: flake8
+49 -128
View File
@@ -4,7 +4,6 @@ import time
from asyncio import AbstractEventLoop
from threading import Thread
import psutil
from aio_pika import Message
from aio_pika.abc import AbstractRobustChannel
@@ -47,47 +46,6 @@ class ScaleRequestV1:
)
class RunningAverage:
"""
A class to maintain a running average of the last N values.
The intended use is to track CPU usage over time window, that is directly related to sample rate and window size.
This is the reason why the constructor takes the number of items to store, calculated as window size / sample rate
Initial_value is for unit test to emulate OVERLOAD or IDLE condition.
"""
def __init__(self, num_items, initial_value=0):
self.buffer = [initial_value] * num_items # Fixed-size array
self.pointer = 0 # Current position in array
self.num_items = num_items # Maximum number of items to store
self.is_filled = False # Track if buffer is fully filled
def add(self, value):
# Store the value at current pointer position
self.buffer[self.pointer] = value
# Move pointer to next position with wrap-around
self.pointer += 1
if self.pointer >= self.num_items:
self.pointer = 0
self.is_filled = True
def get_average(self):
# Determine how many items we should average
count = self.num_items if self.is_filled else self.pointer
if count == 0:
return 0.0 # Avoid division by zero
return sum(self.buffer) / count
def get_current_values(self):
"""Returns the values in chronological order (oldest first)"""
if not self.is_filled:
return self.buffer[: self.pointer]
return self.buffer[self.pointer :] + self.buffer[: self.pointer]
class BackpressureHandler:
# Track the number of messages currently being processed
current_parallel_executions = 0
@@ -96,8 +54,6 @@ 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,
@@ -113,19 +69,6 @@ class BackpressureHandler:
self.swarm_service_id = self.config.amq_adapter.swarm_service_id
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"""
@@ -156,7 +99,6 @@ 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
@@ -176,93 +118,70 @@ class BackpressureHandler:
self.last_backpressure_event_time = time.time()
self.last_backpressure_event = ScalingRequestAlert.OVERLOAD
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.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
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
def backpressure_monitor_loop(self):
_time_window_sec = self.config.backpressure.time_window
_idle_duration_millis = 1000 * self.config.backpressure.idle_duration
_loop = asyncio.new_event_loop()
async def _backpressure_monitor():
"""Monitor the backpressure conditions and trigger events accordingly"""
while True:
logging_info(
"Backpressure: Monitor loop, current=%s",
self.current_parallel_executions,
)
logging_info(
"Backpressure: Last data message time=%s, eventTime=%s",
self.last_data_message_time,
self.last_backpressure_event_time,
)
logging_info(
"Backpressure: Last backpressure event=%s",
self.last_backpressure_event,
)
# 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
)
self.current_parallel_executions >= self.parallel_workers
and time.time() - self.last_backpressure_event_time
> _time_window_sec
and 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_time = time.time()
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
)
self.current_parallel_executions == 0
and time.time() - self.last_data_message_time
>= _idle_duration_millis
and 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
self.last_backpressure_event_time = time.time()
else:
# Trigger update event
if _now - self.last_backpressure_event_time > _time_window_millis:
if (
self.current_parallel_executions > 0
and time.time() - self.last_backpressure_event_time
> _time_window_sec
):
# Trigger the update event
_scaling_request: ScaleRequestV1 = ScaleRequestV1(
self.swarm_service_id,
self.swarm_task_id,
100,
_average_cpu_usage,
self.parallel_workers,
self.parallel_workers - self.current_parallel_executions,
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_time = time.time()
self.last_backpressure_event = ScalingRequestAlert.UPDATE
await asyncio.sleep(_monitor_interval)
await asyncio.sleep(_time_window_sec)
def backpressure_monitor_loop(self):
_loop = asyncio.new_event_loop()
_loop.run_until_complete(self._backpressure_monitor())
_loop.run_until_complete(_backpressure_monitor())
async def handle_backpressure_overload_event(self):
logging_warning("Backpressure: Capacity close to depleted!")
@@ -302,13 +221,15 @@ class BackpressureHandler:
if not self.exchange:
async def _wrap_rabbit_mq_api_init(channel):
_exchange = await channel.get_exchange(name="cleverthis.clevermicro.management")
_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)
asyncio.run_coroutine_threadsafe(
_wrap_rabbit_mq_api_init(self.channel), self.loop
)
)
if self.exchange:
@@ -336,6 +257,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))
await await_future(
asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api(), self.loop)
)
+171 -528
View File
@@ -4,71 +4,53 @@ Methods here will invoke the actual CleverThis Service code, which are passed as
"""
import asyncio
import enum
import inspect
import json
import pathlib
import re
from asyncio import AbstractEventLoop
from collections import defaultdict
from typing import Any, Callable, Coroutine, Dict, List
from typing import Dict, List, Callable, Coroutine, Any
from aio_pika.abc import (
AbstractExchange,
AbstractRobustChannel,
AbstractRobustConnection,
AbstractExchange,
)
from aiohttp import ClientResponse, ClientSession
from fastapi import HTTPException, UploadFile
from fastapi.params import File
from fastapi.routing import APIRoute
from fastapi.security import OAuth2PasswordRequestForm
from pydantic import BaseModel
from starlette.datastructures import Headers
from aiohttp import ClientSession, ClientResponse
from dataclasses import dataclass
from fastapi import HTTPException
from starlette.responses import FileResponse
from amqp.adapter import serializer
from amqp.adapter.data_parser import AMQDataParser
from amqp.adapter.data_response_factory import DataResponseFactory
from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
from amqp.adapter.logging_utils import (
logging_debug,
logging_error,
logging_info,
logging_warning,
)
from amqp.adapter.pydantic_serializer import python_type_to_json
from amqp.adapter.type_descriptor import TypeDescriptor
from amqp.adapter.logging_utils import logging_error, logging_debug, logging_info
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQErrorMessage, AMQMessage, DataMessage, DataResponse
from amqp.model.model import DataMessage, DataResponse
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE
from amqp.router.utils import await_result
g_extra_init_data: dict = dict()
async def _create_reply_exchange_and_queue(
connection: AbstractRobustConnection,
) -> tuple[AbstractRobustChannel, AbstractExchange, str]:
@dataclass
class AMQMessage(DataMessage):
"""
Create objects for reply.
Return (channel for reply queue, exchange for reply, reply queue name)
Expand the core DataMessage type to include the channel, which is required to create a dedicated queue
in case the response should initiate file download.
"""
# 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
def __init__(self, data_message: DataMessage, connection: AbstractRobustConnection):
# Copy all fields from DataMessage
self.magic = data_message.magic
self.id = data_message.id
self.method = data_message.method
self.domain = data_message.domain
self.path = data_message.path
self.headers = data_message.headers
self.trace_info = data_message.trace_info
self.base64_body = data_message.base64_body
# Add the new field
self.connection = connection
self.extra_data = (
data_message.extra_data if hasattr(data_message, "extra_data") else {}
)
logging_debug(f"Publishing to exchange: {_exchange_name}, q: {_queue_name}")
return _channel, _exchange, _queue_name
async def _convert_file_response(
@@ -79,22 +61,41 @@ async def _convert_file_response(
) -> str:
"""
Converts a FileResponse object to a JSON string containing the filename and stream name.
The file content itself is transmitted through the named stream (which is implemented as
temporary dedicated queue.
The file content itself is transmitted through the named stream (which is implemented as temporary dedicated queue.
"""
# All calls to RabbitMQ API have to be called threadsafe, because those appear to be attached
# to parent event loop, and need to execute at that event loop, not the current one
# allow coroutine to execute. calling _future.result() will block the entire thread
# and the coroutine will never execute
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
# the 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(_create_reply_exchange_and_queue(connection), loop)
asyncio.run_coroutine_threadsafe(_wrap_amq_api_calls(connection), loop)
)
# publish_file also needs RabbitMQ API, ensure the loop is passed as argument
(_stream_name, _file_size) = await file_handler.publish_file(
_exchange, obj.path, _routing_key, loop
)
# as above, but luckily this time we don't need to wait for the result
asyncio.run_coroutine_threadsafe(_channel.close(), loop)
_future = asyncio.run_coroutine_threadsafe(_channel.close(), loop)
return json.dumps(
{
"files": [
@@ -111,8 +112,6 @@ async def _convert_file_response(
def is_likely_json(text: str) -> bool:
"""Quick test to see if should convert the response string to JSON format."""
if not isinstance(text, str):
return False
_text = text.strip()
return (_text.startswith("{") and _text.endswith("}")) or (
_text.startswith("[") and _text.endswith("]")
@@ -122,28 +121,20 @@ def is_likely_json(text: str) -> bool:
# ============= CleverThisServiceAdapter =============
class CleverThisServiceAdapter:
"""
A base class for CleverThis service adapter that handles the incoming AMQP messages and routes
them to the appropriate CleverThisService API endpoint.
IMPORTANT: This class is not intended to be used directly. It should be subclassed for
each CleverThis service and override the methods to handle the specific API calls.
A base class for CleverThis service adapter that handles the incoming AMQP messages and routes them to the
appropriate CleverThisService API endpoint.
IMPORTANT: This class is not intended to be used directly. It should be subclassed for each CleverThis service
and override the methods to handle the specific API calls.
"""
def __init__(
self,
amq_configuration: AMQConfiguration,
routes: List[APIRoute],
session: ClientSession = None,
authenticated_user_arg_name: str = "authenticated_user",
extra_init_data=None,
self, amq_configuration: AMQConfiguration, session: ClientSession = None
):
"""
Initializes the CleverThisServiceAdapter with the given AMQ configuration and optional HTTP
session. The HTTP session is for the case when the adapter uses loose coupling with
the CleverThis service via HTTP REST API.
Initializes the CleverThisServiceAdapter with the given AMQ configuration and optional HTTP session.
The HTTP session is for the case when the adapter uses loose coupling with the CleverThis service via
HTTP REST API.
"""
if extra_init_data is None:
extra_init_data = dict()
self.authenticated_user_arg_name = authenticated_user_arg_name
self.amq_configuration = amq_configuration
self.session = session
self.service_host = self.amq_configuration.dispatch.service_host
@@ -152,81 +143,6 @@ class CleverThisServiceAdapter:
self.require_authenticated_user = (
self.amq_configuration.amq_adapter.require_authenticated_user
)
self.routes = self.group_routes_by_method(api_routes=routes)
global g_extra_init_data
g_extra_init_data = extra_init_data or {}
def get_endpoint(self, endpoint: Any) -> Any:
"""
Optionally, Service can override this method to alter the endpoint mapping.
This is useful for example when the CleverThis service has a JSON wrapper
around the endpoint, like CleverSwarm, which might be the preferred endpoint to call.
:param endpoint: The original endpoint to be called.
:return: The actual endpoint to be called.
"""
return endpoint
def clone_and_adapt_route(self, original_route: APIRoute) -> APIRoute:
"""
Creates a clone of APIRoute. Checks if there's preferred endpoint with json serialization.
It also adjusts the path regex to ensure it matches the path in different API versions.
"""
# Need to adapt also the path regex to ensure it matches the path in different API versions
# Ensures the leading '^' is followed by any match and there's no '/' before final '$'
_new_regex_string = "^.*" + original_route.path_regex.pattern.lstrip("^").rstrip("/$") + "$"
_new_route: APIRoute = APIRoute(
path=original_route.path,
endpoint=self.get_endpoint(
original_route.endpoint
), # Service can alter the endpoint mapping here
response_model=original_route.response_model,
status_code=original_route.status_code,
tags=original_route.tags.copy() if original_route.tags else None,
dependencies=(
original_route.dependencies.copy() if original_route.dependencies else None
),
summary=original_route.summary,
description=original_route.description,
response_description=original_route.response_description,
responses=(original_route.responses.copy() if original_route.responses else None),
deprecated=original_route.deprecated,
name=original_route.name,
methods=original_route.methods.copy() if original_route.methods else None,
operation_id=original_route.operation_id,
response_model_include=original_route.response_model_include,
response_model_exclude=original_route.response_model_exclude,
response_model_by_alias=original_route.response_model_by_alias,
response_model_exclude_unset=original_route.response_model_exclude_unset,
response_model_exclude_defaults=original_route.response_model_exclude_defaults,
response_model_exclude_none=original_route.response_model_exclude_none,
include_in_schema=original_route.include_in_schema,
response_class=original_route.response_class,
dependency_overrides_provider=original_route.dependency_overrides_provider,
callbacks=(original_route.callbacks.copy() if original_route.callbacks else None),
openapi_extra=(
original_route.openapi_extra.copy() if original_route.openapi_extra else None
),
generate_unique_id_function=original_route.generate_unique_id_function,
)
_new_route.path_regex = re.compile(_new_regex_string)
return _new_route
def group_routes_by_method(self, api_routes: List[APIRoute]) -> Dict[str, List[APIRoute]]:
"""
Groups API routes by their HTTP method. This is important to ensure that URLPath is matched
correctly, as different methods may share the same URLPath (e.g. GET, POST PUT, DELETE).
"""
_method_to_routes = defaultdict(list) # Automatically initializes new lists for new keys
for _route in api_routes:
if (
isinstance(_route, APIRoute) and _route.methods
): # Ensure methods exist (should always be true for valid routes)
for method in _route.methods:
_method_to_routes[method].append(self.clone_and_adapt_route(_route))
return dict(_method_to_routes)
def get_content_type(self, message_headers: Dict[str, List[str]]) -> str:
"""
@@ -241,11 +157,9 @@ class CleverThisServiceAdapter:
async def get_current_user(self, token: str) -> object:
"""
To Be Overriden in subclass for given CleverThis Service, to return the current user based
on the token, in format appropriate for the service.
To Be Overriden in subclass for given CleverThis Service, to return the current user based on the token,
in format appropriate for the service.
:param token: The token to be used for authentication.
:return: A user object or None if the user is not authenticated.
"""
return None
@@ -259,9 +173,11 @@ class CleverThisServiceAdapter:
_auth_user: Any | None = None
_amq_response: Any | None = None
logging_debug(f"on_message: require_authenticated_user={self.require_authenticated_user}")
logging_debug(
f"on_message: require_authenticated_user={self.require_authenticated_user}"
)
if self.require_authenticated_user:
_auth: str = message.headers().get("Authorization", "")
_auth: str = message.headers.get("Authorization", "")
logging_debug(f"Auth: {_auth}")
if isinstance(_auth, List):
_auth = _auth[0]
@@ -273,8 +189,12 @@ class CleverThisServiceAdapter:
except Exception as error:
logging_error(f"on_message: Error getting user: {error}")
if _auth_user or not self.require_authenticated_user or message.path().endswith("/login"):
method = message.method().lower()
if (
_auth_user
or not self.require_authenticated_user
or message.path.endswith("/login")
):
method = message.method.lower()
if method == "get":
_amq_response = await self.get(message, _auth_user)
elif method == "put":
@@ -286,419 +206,160 @@ class CleverThisServiceAdapter:
elif method == "delete":
_amq_response = await self.delete(message, _auth_user)
else:
raise ValueError(f"Unexpected HTTP method: {message.method()}")
raise ValueError(f"Unexpected HTTP method: {message.method}")
else:
_amq_response = DataResponseFactory.of_error_message(
message.id(),
AMQErrorMessage(401, "Unauthorized", "Unauthorized"),
message.trace_info(),
_amq_response = DataResponseFactory.of(
message.id, 401, "text/plain", b"Unauthorized", message.trace_info
)
logging_info(
f"ALL DONE in on_message id:{message.id()}, resp code:{_amq_response.response_code()}"
f"ALL DONE in on_message id:{message.id}, resp code:{_amq_response.response_code}"
)
return _amq_response
async def _authorize_forward_request(
self, message: DataMessage, request_coroutine, auth_user: Any
) -> DataResponse:
# ==================================================================================================
# ===================== C l e v e r T h i s A P I m e t h o d s ============================
# ==================================================================================================
async def get(self, message: DataMessage, auth_user: Any) -> DataResponse:
"""
This method is used to authorize the request and forward it to the CleverThis Service API.
Used only for loose / 3rd party mode of coupling with the Service
"""
# AMQDataParser.parse_request_body(message)
# TODO: implement authorization logic here - ensure the security related headers are set
# This depends on what kind of security is required by the CleverThis service.
# However at this time there is no loosely coupled Service.
return await request_coroutine
# ==============================================================================================
# ===================== C l e v e r T h i s A P I m e t h o d s ========================
# ==============================================================================================
async def get(self, message: AMQMessage, auth_user: Any) -> DataResponse:
"""
Handles the GET requests for the CleverThis API. In tight coupling (self.session is None),
use path, path params and query string to determine the endpoint method to invoke.
In loose coupling (self.session is not None), forward the request as REST request using
the session, i.e. attempt to call the REST endpoint directly using HTTP protocol.
Handles the GET requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
It will attempt to call the REST endpoint directly using HTTP protocol.
:param message: message from the AMQP that contains the GET request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
if self.session is None:
return await self.handle_no_body_payload(message, auth_user, "GET")
# Forwards as REST - try to invoke the service directly on this path
url = f"http://{self.service_host}:{self.service_port}{message.path()}"
headers = {**message.headers(), **message.trace_info()}
# unmatched path - try to invoke the service directly on this path
url = f"http://{self.service_host}:{self.service_port}{message.path}"
headers = {**message.headers, **message.trace_info}
logging_info(
f"SEND HEADERS out: {', '.join(f'{k}={v}' for k, v in headers.items())}, REQ={url}"
)
if self.session:
_http_resp: ClientResponse = await self.session.get(url, headers=headers)
return DataResponseFactory.of(
message.id(),
message.id,
_http_resp.status,
_http_resp.content_type,
await _http_resp.read(),
message.trace_info(),
message.trace_info,
)
return DataResponseFactory.of(
message.id,
404,
"text/plain",
str("Destination location does not exists").encode("utf-8"),
message.trace_info,
)
async def put(self, message: AMQMessage, auth_user: Any) -> DataResponse:
async def put(self, message: DataMessage, auth_user: Any) -> DataResponse:
"""
Handles the PUT requests for the CleverThis API. Uses path, path params, query string and
body payload to determine the endpoint method to invoke and its input arguments
In loose coupling (self.session is not None), attempt to call the REST endpoint directly
using HTTP protocol.
Handles the PUT requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
It will attempt to call the REST endpoint directly using HTTP protocol.
:param message: message from the AMQP that contains the PUT request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
if self.session is None:
return await self.handle_body_payload(message, auth_user, "PUT")
# Forwards as REST - try to invoke the service directly on this path
return await self._authorize_forward_request(
return await self.handle_possible_form(
message,
self.session.put(
f"http://{self.service_host}:{self.service_port}{message.path()}",
data=message.body(),
headers={**message.headers(), **message.trace_info()},
f"http://{self.service_host}:{self.service_port}{message.path}",
data=message.body,
headers={**message.headers, **message.trace_info},
),
auth_user,
)
async def post(self, message: AMQMessage, auth_user: Any) -> DataResponse:
async def post(self, message: DataMessage, auth_user: Any) -> DataResponse:
"""
Handles the POST requests for the CleverThis API. Uses path, path params, query string and
body payload to determine the endpoint method to invoke and its input arguments
In loose coupling (self.session is not None), attempt to call the REST endpoint directly
using HTTP protocol.
:param message: message from AMQP that contains the POST request, including body and headers
Handles the POST requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
It will attempt to call the REST endpoint directly using HTTP protocol.
:param message: message from the AMQP that contains the POST request, including body and headers
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
if self.session is None:
return await self.handle_body_payload(message, auth_user, "POST")
return await self._authorize_forward_request(
return await self.handle_possible_form(
message,
self.session.post(
f"http://{self.service_host}:{self.service_port}{message.path()}",
data=message.body(),
headers={**message.headers(), **message.trace_info()},
f"http://{self.service_host}:{self.service_port}{message.path}",
data=message.body,
headers={**message.headers, **message.trace_info},
),
auth_user,
)
async def head(self, message: AMQMessage, auth_user: Any) -> DataResponse:
async def handle_possible_form(
self, message: DataMessage, request_coroutine, auth_user: Any
) -> DataResponse:
_args = AMQDataParser.parse_request_body(message)
return await request_coroutine
async def head(self, message: DataMessage, auth_user: Any) -> DataResponse:
"""
Handles the HEAD requests for the CleverThis API. In tight coupling (self.session is None),
use path, path params and query string to determine the endpoint method to invoke.
In loose coupling (self.session is not None), HEAD call is not supported.
:param message: message from AMQP containing the HEAD request, including body and headers
:param auth_user: user authenticated by the API Gateway, as UserSchema
Handles the HEAD requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
It will attempt to call the REST endpoint directly using HTTP protocol.
:param message: message from the AMQP that contains the POST request, including body and headers
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
if self.session is None:
return await self.handle_no_body_payload(message, auth_user, "HEAD")
raise NotImplementedError
async def delete(self, message: AMQMessage, auth_user: Any) -> DataResponse:
async def delete(self, message: DataMessage, auth_user: Any) -> DataResponse:
"""
Handles the DELETE requests for the CleverThis API.
:param message: message from AMQP that contains the DELETE request, including body & headers
:param auth_user: user authenticated by the API Gateway, as UserSchema
Handles the DELETE requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
It will attempt to call the REST endpoint directly using HTTP protocol.
:param message: message from the AMQP that contains the POST request, including body and headers
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
if self.session is None:
return await self.handle_no_body_payload(message, auth_user, "DELETE")
return await self._authorize_forward_request(
message,
self.session.post(
f"http://{self.service_host}:{self.service_port}{message.path()}",
data=message.body(),
headers={**message.headers(), **message.trace_info()},
),
auth_user,
)
async def options(self, message: AMQMessage, auth_user: Any) -> DataResponse:
"""
Handles the OPTIONS requests. In tight coupling (self.session is None),
use path, path params and query string to determine the endpoint method to invoke.
In loose coupling (self.session is not None), forward the request as REST request using
the session, i.e. attempt to call the REST endpoint directly using HTTP protocol.
:param message: message from AMQP that contains the OPTIONS request, including headers
:param auth_user: user authenticated by the API Gateway, as UserSchema
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
if self.session is None:
return await self.handle_no_body_payload(message, auth_user, "OPTIONS")
raise NotImplementedError
async def patch(self, message: AMQMessage, auth_user: Any) -> DataResponse:
"""
Handles the PATCH requests for the CleverBRAG API.
:param message: message from AMQP that contains the PATCH request, including body & headers
:param auth_user: user authenticated by the API Gateway, as UserSchema
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
if self.session is None:
return await self.handle_no_body_payload(message, auth_user, "PATCH")
return await self._authorize_forward_request(
message,
self.session.post(
f"http://{self.service_host}:{self.service_port}{message.path()}",
data=message.body(),
headers={**message.headers(), **message.trace_info()},
),
auth_user,
)
# ==============================================================================================
# ================= C l e v e r S w a r m p a y l o a d h a n d l e r s ==================
# ==============================================================================================
async def handle_no_body_payload(
self, message: AMQMessage, auth_user: Any, method: str
) -> DataResponse:
"""
Handles the requests that do not have a body payload.
:param message: message from the AMQP that contains the request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
_routes = self.routes.get(method, [])
_path, _args = AMQDataParser.parse_get_url(message.path())
for _route in _routes:
_match = _route.path_regex.match(_path)
if _match:
_args.update(_match.groupdict())
_signature_args = inspect.signature(_route.endpoint).parameters
if self.authenticated_user_arg_name in _signature_args.keys():
_args[self.authenticated_user_arg_name] = auth_user
# If the path matches, call the corresponding function
return await self.call_cleverthis_api(
message,
_args,
api_method=_route.endpoint,
success_code=_route.status_code,
)
return DataResponseFactory.of(
message.id(),
404,
"text/plain",
str(f"Destination location [{message.path()}] does not exists").encode("utf-8"),
message.trace_info(),
)
async def handle_body_payload(
self, message: AMQMessage, auth_user: Any, method: str
) -> DataResponse:
"""
Handles the requests that have a body payload, including Multipart data.
:param message: message from the AMQP that contains the request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:param method: HTTP method (POST, PUT, etc.)
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
_form_data = AMQDataParser.parse_request_body(message)
_routes = self.routes.get(method, [])
_path, _query_args = AMQDataParser.parse_get_url(message.path())
for _route in _routes:
_match = _route.path_regex.match(_path)
if _match:
# If the path matches, call the corresponding function
_form_data.update(_match.groupdict())
_form_data.update(_query_args)
for _body_param in _route.dependant.body_params:
if _body_param.name in _form_data:
_file_data = _form_data.get(_body_param.name, None)
actual_filepath = message.extra_data.get(_body_param.name, None)
_form_data[_body_param.name] = (
(
_body_param.type_.__name__ == "UploadFile"
or isinstance(_body_param.field_info, File)
)
and _file_data is not None
and actual_filepath is not None
and UploadFile(
file=pathlib.Path(actual_filepath).open("rb"),
size=_file_data.get("size"),
filename=_file_data.get("filename"),
headers=Headers({"Content-Type": _file_data.get("mediaType")}),
)
or _form_data[_body_param.name]
)
_verified_args = {}
_required_args = getattr(_route.endpoint, "__annotations__", {})
_signature_args = inspect.signature(_route.endpoint).parameters
for arg in _required_args:
if arg in _form_data:
_verified_args[arg] = _form_data[arg]
else:
# needed for CleverSwarm /login input
if _required_args[arg].__name__ == "OAuth2PasswordRequestForm":
try:
_verified_args[arg] = OAuth2PasswordRequestForm(**_form_data)
except Exception as e:
print(f"Error: {e}")
if self.authenticated_user_arg_name in _signature_args.keys():
_verified_args[self.authenticated_user_arg_name] = auth_user
return await self.call_cleverthis_api(
message,
_verified_args,
api_method=_route.endpoint,
success_code=_route.status_code if _route.status_code else 200,
)
return DataResponseFactory.of(
message.id(),
404,
"text/plain",
str(f"Destination location [{message.path()}] does not exists").encode("utf-8"),
message.trace_info(),
)
def _verify_and_instantiate_arguments(
self,
@staticmethod
async def call_cleverthis_api(
message: AMQMessage,
args: Any,
api_method: Callable[[Any, Any], Coroutine],
) -> dict:
success_code: int,
loop: AbstractEventLoop | None = None,
) -> DataResponse:
"""
Check the provided arguments against the API method signature and convert the incoming string value
into correct expected input object.
The logic used FastAPI annotation and would use the value from 'default_factory' invocation (where defined)
or 'default' value, as defined in the FastAPI annotation, to create default value in case the object is not in
the input values provided by the caller.
Handles POST and PUT requests, which typically contain possibly complex payload data, passed as a dictionary.
Invokes the provided API method with the given arguments and authenticated user,
and returns the response as DataResponse.
:param message: DataMessage containing the details of the call and the payload
:param args: map of the arguments to pass to the API method
:param api_method: Callable - the method to invoke. Last parameter is the authenticated user
:param success_code: HTTP code to be returned when operation suceeds
:return: DataResponse class
"""
_verified_args = {}
try:
_verified_args = {}
_sig = inspect.signature(api_method)
_required_args = list(_sig.parameters.values())
# _required_args = getattr(api_method, '__annotations__', {})
for arg in _required_args:
if arg.name in args:
_annotation = arg.annotation
_str_annotation = str(_annotation)
# convert to another type if arg is not a str
if arg.name == self.authenticated_user_arg_name:
_verified_args[arg.name] = args.get(arg.name)
elif _annotation == int:
_verified_args[arg.name] = int(args.get(arg.name))
elif _annotation == float:
_verified_args[arg.name] = float(args.get(arg.name))
elif _annotation == bool:
_verified_args[arg.name] = serializer.str_to_bool(args.get(arg.name))
elif (
"Optional" in _str_annotation
or _str_annotation.lower().startswith("list")
or _str_annotation.startswith("typing.List")
or "class" in _str_annotation
):
_type_descr = TypeDescriptor(
_str_annotation, extra_init_data=g_extra_init_data
)
try:
_payload = args.get(arg.name, "")
_verified_args[arg.name] = _type_descr.instantiate(
_payload, is_json=is_likely_json(_payload)
)
except Exception as err:
logging_error(
f"Could not instantiate argument {arg}, defaulting to its string value",
err,
)
_verified_args[arg.name] = args.get(arg.name, None)
if arg.annotation == int:
_verified_args[arg.name] = int(args[arg.name])
elif arg.annotation == float:
_verified_args[arg.name] = float(args[arg.name])
elif arg.annotation == bool:
_verified_args[arg.name] = bool(args[arg.name])
# TODO: more?
else:
_temp = args.get(arg.name)
if isinstance(arg.annotation, enum.EnumType) or isinstance(
arg.annotation, enum.Enum
):
try:
_temp = arg.annotation[_temp]
except KeyError as ke:
# some enums have string value different to key spelling, try different thing
try:
_temp = arg.annotation(_temp)
except ValueError as ve:
logging_warning(
f"cannot convert '{_temp}' to Enum type {arg.annotation}, {ke}, {ve}"
)
print(f"ENUM -> {_temp}")
_verified_args[arg.name] = _temp
_verified_args[arg.name] = args[arg.name]
else:
if hasattr(arg, "default"):
_temp = arg.default
if hasattr(_temp, "default_factory") and _temp.default_factory is not None:
_temp = _temp.default_factory()
elif hasattr(_temp, "default"):
_temp = _temp.default
if not str(_temp) == "<class 'inspect._empty'>":
_verified_args[arg.name] = _temp
else:
logging_warning(
logging_info(
f"{api_method.__name__}: Missing required argument: {arg.name}"
)
except Exception as e:
logging_error("Failed to validate and instantiate input args", e)
return _verified_args
async def call_cleverthis_api(
self,
message: AMQMessage,
args: Any,
api_method: Callable[[Any, Any], Coroutine],
success_code: int,
) -> DataResponse:
"""
Handles POST and PUT requests, which typically contain possibly complex payload data,
passed as a dictionary. Invokes the provided API method with the given arguments and
authenticated user, and returns the response as DataResponse.
:param message: DataMessage containing the details of the call and the payload
:param args: map of the arguments to pass to the API method
:param api_method: Callable - the method to invoke. Last parameter is the authenticated user
:param success_code: HTTP code to be returned when operation suceeds
:return: DataResponse class
"""
try:
_verified_args = self._verify_and_instantiate_arguments(
args=args, api_method=api_method
)
logging_info(
"INVOKE ENDPOINT: [%s], ARGS: {%s}",
api_method.__name__,
", ".join(f"{k}={v}" for k, v in _verified_args.items()),
f"INVOKE ENDPOINT: {api_method.__name__}, ARGS: {', '.join(f'{k}={v}' for k, v in _verified_args.items())}"
)
_result: Any = await api_method(**_verified_args)
logging_info(f"ENDPOINT RESULT: [{api_method.__name__}], RESULT: {str(_result)}")
logging_info(
f"ENDPOINT RESULT: {api_method.__name__}, RESULT: {str(_result)}"
)
if isinstance(_result, FileResponse):
_json_result, _content_type = (
@@ -706,18 +367,12 @@ class CleverThisServiceAdapter:
_result,
message.connection,
file_handler=StreamingFileHandler(),
loop=self.loop,
loop=loop,
),
"application/octet-stream",
)
elif isinstance(_result, BaseModel):
_json_result = _result.json()
_content_type = "application/json"
else:
if isinstance(_result, str):
_json_result = _result
else:
_json_result = python_type_to_json(_result)
_content_type = "application/json"
_binary_result = (
@@ -726,56 +381,44 @@ class CleverThisServiceAdapter:
else (
_json_result.body
if hasattr(_json_result, "body")
else (_json_result.read() if hasattr(_json_result, "read") else _json_result)
else (
_json_result.read()
if hasattr(_json_result, "read")
else _json_result
)
)
)
logging_info(
f"ALL DONE in call_cleverthis_api id:{message.id()}, resp code:{success_code}"
f"ALL DONE in call_cleverthis_api id:{message.id}, resp code:{success_code}"
)
return DataResponseFactory.of(
message.id(),
message.id,
success_code,
_content_type,
_binary_result,
message.trace_info(),
message.trace_info,
)
except HTTPException as http_error:
logging_error(f"Service REST endpoint invocation failed: {http_error}")
_content = str(http_error.detail)
if not is_likely_json(_content):
return DataResponseFactory.of_error_message(
message.id(),
AMQErrorMessage(
http_error.status_code,
"Service Invocation Failed",
_content,
),
message.trace_info(),
)
_content = '{"error": "' + _content + '"}'
logging_error(f"Service REST endpoint invocation failed: {http_error}")
return DataResponseFactory.of(
message.id(),
message.id,
http_error.status_code,
"application/json",
_content.encode("utf-8"),
message.trace_info(),
message.trace_info,
)
except Exception as error:
logging_error(f"Service endpoint invocation failed: {error}")
_content = str(error)
if not is_likely_json(_content):
return DataResponseFactory.of_error_message(
message.id(),
AMQErrorMessage(
500,
"Service Invocation Failed",
_content,
),
message.trace_info(),
)
_content = '{"error": "' + _content + '"}'
return DataResponseFactory.of(
message.id(),
message.id,
500,
"application/json",
_content.encode("utf-8"),
message.trace_info(),
message.trace_info,
)
+28 -23
View File
@@ -5,21 +5,21 @@ methods to create, serialize and deserialize DataMessage objects.
import base64
import json
import re
from datetime import datetime, timezone
from io import BytesIO
from datetime import datetime, timezone
from typing import Dict, List
import amqp.adapter.file_handler
from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
from amqp.model.model import DataMessage
from amqp.model.snowflake_id import SnowflakeId
from amqp.adapter.serializer import (
map_as_string,
map_with_list_as_string,
map_as_string,
parse_map_list_string,
parse_map_string,
)
from amqp.model.model import DataMessage, DataMessageMagicByte
from amqp.model.snowflake_id import SnowflakeId
from amqp.router.utils import sanitize_as_rabbitmq_name
from amqp.adapter.file_handler import StreamingFileHandler, FileHandler
class DataMessageFactory:
@@ -30,6 +30,7 @@ class DataMessageFactory:
MACHINE_ID_MASK = (1 << MACHINE_ID_BITS) - 1
SEQUENCE_MASK = (1 << SEQUENCE_BITS) - 1
MAGIC_BYTE_HTTP_REQUEST = "A"
_instance = None
snowflake_epoch = datetime(2024, 10, 1, 0, 0, tzinfo=timezone.utc).timestamp()
@@ -99,10 +100,12 @@ class DataMessageFactory:
:return: DataMessage object
"""
serializable_headers = headers.copy() if headers else {}
payload = base64.b64encode(message.encode("utf-8") if isinstance(message, str) else message)
payload = base64.b64encode(
message.encode("utf-8") if isinstance(message, str) else message
)
trace_info = {}
return DataMessage(
DataMessageMagicByte.HTTP_REQUEST.value,
self.MAGIC_BYTE_HTTP_REQUEST,
self.generate_snowflake_id(),
method,
domain,
@@ -119,7 +122,7 @@ class DataMessageFactory:
:param message: DataMessage input
:return: routing key (as string) compliant to RabbitMQ allowed character set.
"""
return sanitize_as_rabbitmq_name(f"{message.domain()}.{message.path()}")
return sanitize_as_rabbitmq_name(f"{message.domain}.{message.path}")
@staticmethod
def get_timestamp_from_id(_id: SnowflakeId) -> int:
@@ -129,7 +132,9 @@ class DataMessageFactory:
:return: the message timestamp
"""
bits = _id.get_bits()
low = bits[0] >> (DataMessageFactory.SEQUENCE_BITS + DataMessageFactory.MACHINE_ID_BITS)
low = bits[0] >> (
DataMessageFactory.SEQUENCE_BITS + DataMessageFactory.MACHINE_ID_BITS
)
high = bits[1] << (
64 - DataMessageFactory.SEQUENCE_BITS - DataMessageFactory.MACHINE_ID_BITS
)
@@ -144,13 +149,13 @@ class DataMessageFactory:
"""
return (
f"DataMessage{{"
f"id={message.id()}, "
f"method='{message.method()}', "
f"domain='{message.domain()}', "
f"path='{message.path()}', "
f"headers={{{map_with_list_as_string(message.headers())}}}, "
f"traceInfo={{{map_as_string(message.trace_info())}}}, "
f"base64body={message.base64body().decode()}"
f"id={message.id}, "
f"method='{message.method}', "
f"domain='{message.domain}', "
f"path='{message.path}', "
f"headers={{{map_with_list_as_string(message.headers)}}}, "
f"traceInfo={{{map_as_string(message.trace_info)}}}, "
f"base64body={message.base64_body.decode()}"
f"}}"
)
@@ -161,19 +166,19 @@ class DataMessageFactory:
:param message: message to serialize
:return: messages as bytes
"""
id_bytes = str(message.id()).encode()
id_bytes = str(message.id).encode()
prolog = (
f"|m={message.method()}|d={message.domain()}|p={message.path()}|h={{{map_with_list_as_string(message.headers())}}}|t={{{map_as_string(message.trace_info())}}}|b="
f"|m={message.method}|d={message.domain}|p={message.path}|h={{{map_with_list_as_string(message.headers)}}}|t={{{map_as_string(message.trace_info)}}}|b="
).encode("utf-8")
header_length = 1 + len(id_bytes) + len(prolog)
prolog_len = header_length.to_bytes(2, byteorder="big")
msg_length = 2 + header_length + len(message.base64body())
msg_length = 2 + header_length + len(message.base64_body)
with BytesIO(bytearray(msg_length)) as baos:
baos.write(prolog_len)
baos.write(DataMessageMagicByte.HTTP_REQUEST.value.encode("utf-8"))
baos.write(DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST.encode("utf-8"))
baos.write(id_bytes)
baos.write(prolog)
baos.write(message.base64body())
baos.write(message.base64_body)
return baos.getvalue()
@staticmethod
@@ -187,7 +192,7 @@ class DataMessageFactory:
metadata = input_bytes[2 : prolog_len + 2].decode()
pattern = re.compile(
f"{DataMessageMagicByte.HTTP_REQUEST.value}"
f"{DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST}"
r"([0-9A-Fa-f]+)\|"
r"m=([^|]*)\|"
r"d=([^|]*)\|"
@@ -215,7 +220,7 @@ class DataMessageFactory:
trace_info = parse_map_string(trace_info_str)
return DataMessage(
DataMessageMagicByte.HTTP_REQUEST.value,
DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST,
SnowflakeId.from_hex(id_str),
method,
domain,
+29 -22
View File
@@ -1,15 +1,15 @@
import io
import json
from typing import Dict, Tuple
import python_multipart
from fastapi import UploadFile
from python_multipart.multipart import Field, File
from typing import Tuple, Dict
from urllib.parse import parse_qs, urlparse
from xml.etree import ElementTree
import python_multipart
from fastapi import UploadFile
from python_multipart.multipart import Field, File
from amqp.adapter.logging_utils import logging_debug, logging_error
from amqp.model.model import AMQMessage
from amqp.adapter.logging_utils import logging_error, logging_debug
from amqp.model.model import DataMessage
class MultipartFormDataParser:
@@ -22,14 +22,15 @@ class MultipartFormDataParser:
included, the message is transmitted in the unchanged multipart/form-data format.
"""
def __init__(self, message: AMQMessage):
def __init__(self, message: DataMessage):
self.form_data: dict = {}
self.is_multipart = False
self.is_json = False
self.is_valid = True
# Convert each list of strings to a single string joined by newline, then to bytes
_headers = {
key: "\n".join(value).encode("utf-8") for key, value in message.headers().items()
key: "\n".join(value).encode("utf-8")
for key, value in message.headers.items()
}
try:
if (
@@ -40,7 +41,7 @@ class MultipartFormDataParser:
try:
self.form_data = json.loads(message.body())
self.is_json = True
if self.form_data["files"] is not None and self.form_data["form"] is not None:
if self.form_data["files"] and self.form_data["form"]:
self.form_data = AMQDataParser.process_form_data(self.form_data)
except Exception as jsonerror:
logging_error(f"Parsing JSON: {jsonerror}")
@@ -51,14 +52,15 @@ class MultipartFormDataParser:
)
self.is_multipart = True
except Exception as e:
__error = f"Error parsing multipart/form-data: {e}. Trying parsing as JSON."
logging_error(
f"Error parsing multipart/form-data: {e}. Trying parsing as JSON."
)
try:
self.form_data = json.loads(message.body())
self.is_json = True
if self.form_data["files"] is not None and self.form_data["form"] is not None:
if self.form_data["files"] and self.form_data["form"]:
self.form_data = AMQDataParser.process_form_data(self.form_data)
except Exception as jsonerror:
logging_error(__error)
logging_error(f"Parsing JSON: {jsonerror}")
self.is_valid = False
@@ -68,9 +70,8 @@ class MultipartFormDataParser:
def on_file(self, file: File):
logging_debug(str(file))
file.file_object.seek(0)
self.form_data[file.field_name.decode("utf-8")] = UploadFile(
file.file_object, size=file.size, filename=file.file_name.decode("utf-8")
file.file_object, size=file.size, filename=file.file_name
)
@@ -110,6 +111,7 @@ 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
@@ -128,11 +130,13 @@ class AMQDataParser:
_parsed = urlparse(url)
_query_params = parse_qs(_parsed.query)
# Convert values from lists to single values when there's only one value
_simplified_params = {k: v[0] if len(v) == 1 else v for k, v in _query_params.items()}
_simplified_params = {
k: v[0] if len(v) == 1 else v for k, v in _query_params.items()
}
return _parsed.path, _simplified_params
@staticmethod
def parse_request_body(message: AMQMessage):
def parse_request_body(message: DataMessage):
"""
Parses the HTTP POST request body based on the MIME type.
This is typically called fom the concrete CleverThis Service to extract parameters from the message body.
@@ -142,14 +146,14 @@ class AMQDataParser:
Returns:
dict: Parsed data as a dictionary.
"""
_content_type = message.headers().get("Content-Type", "")
_content_type = message.headers.get("Content-Type", "")
if not _content_type:
raise ValueError("Content-Type header is required")
content_type = _content_type[0]
# Parse JSON
if content_type == "application/json":
try:
return json.loads(message.body() or "{}")
return json.loads(message.body())
except json.JSONDecodeError as e:
logging_error(f"Invalid JSON: {e}")
raise ValueError(f"Invalid JSON: {e}")
@@ -158,12 +162,15 @@ class AMQDataParser:
elif content_type == "application/x-www-form-urlencoded":
try:
_form_data: dict = {
k: v[0] if len(v) == 1 else v for k, v in parse_qs(message.body()).items()
k: v[0] if len(v) == 1 else v
for k, v in parse_qs(message.body()).items()
}
if len(_form_data) == 0 and len(message.body()) > 0:
_form_data = json.loads(message.body())
if _form_data["files"] is not None and _form_data["form"] is not None:
_form_data = AMQDataParser.process_form_data(_form_data)
if _form_data["files"] and _form_data["form"]:
_form_data = MultipartFormDataParser.process_form_data(
_form_data
)
return _form_data
except Exception as e:
logging_error(f"Invalid URL-encoded form data: {e}")
+18 -44
View File
@@ -5,13 +5,13 @@ method to create, serialize and deserialize DataResponse message
import base64
import re
from io import BytesIO
from amqp.adapter.serializer import map_as_string, parse_map_string
from amqp.model.model import AMQErrorMessage, DataMessageMagicByte, DataResponse
from amqp.model.model import DataResponse
from amqp.model.snowflake_id import SnowflakeId
from amqp.adapter.serializer import parse_map_string, map_as_string
class DataResponseFactory:
MAGIC_BYTE_HTTP_REQUEST = "A"
@staticmethod
def create_async_response_message(_id: SnowflakeId) -> DataResponse:
@@ -20,14 +20,8 @@ class DataResponseFactory:
:param _id: ID
:return: Response message
"""
return DataResponseFactory.of(
id=_id,
response_code=200,
content_type="application/text",
body="OK".encode("utf-8"),
trace_info={},
error="",
error_cause="",
return DataResponse(
str(_id), 200, "application/text", "OK".encode("utf-8"), "", "", {}
)
@staticmethod
@@ -37,8 +31,6 @@ class DataResponseFactory:
content_type: str,
body: bytes,
trace_info: dict[str, str],
error: str | None = None,
error_cause: str | None = None,
) -> DataResponse:
"""
Create the DataResponse message supplying all key values
@@ -50,31 +42,13 @@ class DataResponseFactory:
:return: DataResponseMessage object
"""
return DataResponse(
magic=DataMessageMagicByte.HTTP_REQUEST.value,
id=id,
response_code=str(response_code),
id=id.as_string(),
response_code=response_code,
content_type=content_type,
error=error,
error_cause=error_cause,
headers={},
response=body,
error=None,
error_cause=None,
trace_info=trace_info,
base64body=base64.b64encode(body),
)
@staticmethod
def of_error_message(
id: SnowflakeId,
error_message: AMQErrorMessage,
trace_info: dict[str, str],
) -> DataResponse:
return DataResponseFactory.of(
id=id,
response_code=error_message.response_code,
content_type=error_message.content_type,
body=error_message.serialize(),
trace_info=trace_info,
error=error_message.error,
error_cause=error_message.detail,
)
@staticmethod
@@ -84,19 +58,19 @@ class DataResponseFactory:
:param response: Object to serialize
:return: byte array as serialized message
"""
id_bytes = response.id().as_string().encode("utf-8")
id_bytes = response.id.encode("utf-8")
prolog = (
f"|r={response.response_code()}|c={response.content_type()}|e={response.error()}|f={response.error_cause()}|t={{{map_as_string(response.trace_info())}}}|b="
f"|r={response.response_code}|c={response.content_type}|e={response.error}|f={response.error_cause}|t={{{map_as_string(response.trace_info)}}}|b="
).encode("utf-8")
header_length = 1 + len(id_bytes) + len(prolog)
prolog_len = header_length.to_bytes(2, byteorder="big")
msg_length = 2 + header_length + len(response.base64body())
msg_length = 2 + header_length + len(response.response)
with BytesIO(bytearray(msg_length)) as baos:
baos.write(prolog_len)
baos.write(DataMessageMagicByte.HTTP_REQUEST.value.encode("utf-8"))
baos.write(DataResponseFactory.MAGIC_BYTE_HTTP_REQUEST.encode("utf-8"))
baos.write(id_bytes)
baos.write(prolog)
baos.write(response.base64body())
baos.write(base64.b64encode(response.response))
return baos.getvalue()
@staticmethod
@@ -110,7 +84,7 @@ class DataResponseFactory:
metadata = input_bytes[2 : prolog_len + 2].decode()
pattern = re.compile(
f"{DataMessageMagicByte.HTTP_REQUEST.value}"
f"{DataResponseFactory.MAGIC_BYTE_HTTP_REQUEST}"
r"([0-9A-Fa-f]+)\|"
r"r=([^|]*)\|"
r"c=([^|]*)\|"
@@ -136,14 +110,14 @@ class DataResponseFactory:
trace_info = parse_map_string(trace_info_str)
return DataResponseFactory.of(
return DataResponse(
id_str,
response_code,
content_type,
base64.b64decode(body_b64),
trace_info,
error,
error_cause,
trace_info,
)
+69 -75
View File
@@ -4,22 +4,22 @@ implements the file downloader for the amqp adapter, where file is sent in chunk
import asyncio
import json
import logging
import os
from asyncio import AbstractEventLoop
import logging
from asyncio import Future, AbstractEventLoop
from collections import defaultdict
from pathlib import Path
from aio_pika import Message, connect_robust
from aio_pika import connect_robust, Message
from aio_pika.abc import (
AbstractExchange,
AbstractIncomingMessage,
AbstractQueue,
AbstractRobustChannel,
AbstractQueue,
AbstractExchange,
DeliveryMode,
)
from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
from amqp.adapter.logging_utils import logging_error, logging_debug, logging_info
from amqp.router.utils import await_result
# Global dictionary to track completed downloads
@@ -153,44 +153,6 @@ 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):
@@ -204,6 +166,7 @@ class StreamingFileHandler(FileHandler):
message_data (dict): The entire message
output_dir (str): The directory where the file should be saved.
"""
global download_locations
filename = file_info["filename"]
size = file_info["size"]
stream_name = file_info["streamName"]
@@ -222,31 +185,64 @@ 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(
self._download_from_queue_to_file(rabbitmq_url, queue_name, filepath, loop),
loop=loop,
)
asyncio.run_coroutine_threadsafe(wrap_rabbit_mq_calls(), loop=loop)
)
return _local_file_size, _local_eof
return 0, CANCELLED
async def _download_from_queue_to_buffer(
self, rabbitmq_url, queue_name, loop
) -> tuple[bytearray, int]:
async def download_buffer(self, loop, rabbitmq_url, queue_name) -> (bytearray, int):
"""
Internal method for downloading a file from a queue to RAM.
Return (bytearray, eof status).
Asynchronously downloads a bytearay from RabbitMQ based on the provided file information.
It assumes entire array is in the first (and only) chunk
Args:
loop: The asyncio event loop.
rabbitmq_url: The RabbitMQ connection URL.
queue_name (str): The directory where the file should be saved.
"""
_connection = await self._get_connection(rabbitmq_url, loop=loop)
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}")
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():
@@ -261,21 +257,8 @@ class StreamingFileHandler(FileHandler):
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
Args:
loop: The asyncio event loop.
rabbitmq_url: The RabbitMQ connection URL.
queue_name (str): The directory where the file should be saved.
"""
logging_debug(f"Downloading buffer from stream: {queue_name}")
if queue_name:
_local_res, _local_eof = await await_result(
asyncio.run_coroutine_threadsafe(
self._download_from_queue_to_buffer(rabbitmq_url, queue_name, loop), loop=loop
)
asyncio.run_coroutine_threadsafe(wrap_rabbit_mq_calls(), loop=loop)
)
return _local_res, _local_eof
return 0, CANCELLED
@@ -298,6 +281,7 @@ class StreamingFileHandler(FileHandler):
loop: The asyncio event loop of the RabbitMQ API
output_dir
"""
global download_locations
try:
_message_data = json.loads(message.body.decode())
_amq_message_data = json.loads(json_data.decode())
@@ -314,16 +298,22 @@ class StreamingFileHandler(FileHandler):
files_to_download.append(files_info)
if not files_to_download:
logging_info(" [%s] No files to download in this message.", message.delivery_tag)
logging_info(
" [%s] No files to download in this message.", message.delivery_tag
)
# await message.ack()
return defaultdict()
# Create a task for each file download
_tasks = []
for file_info in files_to_download:
logging_info(f" [{message.delivery_tag}] about to create task for {file_info}")
logging_info(
f" [{message.delivery_tag}] about to create task for {file_info}"
)
_task = asyncio.create_task(
self.download_file(loop, rabbitmq_url, file_info, _message_data, output_dir)
self.download_file(
loop, rabbitmq_url, file_info, _message_data, output_dir
)
)
_tasks.append(_task)
logging_info(
@@ -389,7 +379,9 @@ class StreamingFileHandler(FileHandler):
"total_size": _file_size,
"chunk_size": len(_chunk),
"eof": (
END_OF_FILE if max_chunk_size + _offset >= _file_size else IN_PROGRESS
END_OF_FILE
if max_chunk_size + _offset >= _file_size
else IN_PROGRESS
),
},
)
@@ -402,7 +394,9 @@ class StreamingFileHandler(FileHandler):
while not _future.done():
# sleep allows coroutine to execute. must wait for completion to ensure all chunks written in right order
await asyncio.sleep(0.02)
logging_debug(f"Waiting for chunk {_offset}/{_file_size} future= {_future}")
logging_debug(
f"Waiting for chunk {_offset}/{_file_size} future= {_future}"
)
_offset += len(_chunk)
logging_debug(
+15 -18
View File
@@ -4,11 +4,8 @@ import os
import sys
import threading
import traceback
from importlib.metadata import version
# Get version from installed package metadata
__version__ = version("amq_adapter") # Name must match pyproject.toml [project] name
__verbose__ = os.environ.get("AMQ_VERBOSE_LOGGING", "0") == "1"
verbose = os.environ.get("AMQ_VERBOSE_LOGGING", "0") == "1"
def get_context_info():
@@ -54,18 +51,18 @@ def logging_error(msg: str, *args) -> None:
if _exec_info:
_tb = traceback.extract_tb(sys.exc_info()[2])[-1]
logging.error(
f"{_thread_id}:{__version__} [E] '{msg}' ---> '{_tb.filename}':{_tb.lineno}, '{_tb.name}()'",
f"{_thread_id} [E] '{msg}' ---> '{_tb.filename}':{_tb.lineno}, '{_tb.name}()'",
*args,
)
else:
module, line_number, function_name = get_context_info()
if module and line_number and function_name:
logging.error(
f"{_thread_id}:{__version__} [E] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
f"{_thread_id} [E] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
*args,
)
else:
logging.error(f"{_thread_id}:{__version__} [E] {msg}", *args)
logging.error(f"{_thread_id} [E] {msg}", *args)
def logging_warning(msg: str, *args) -> None:
@@ -76,17 +73,17 @@ def logging_warning(msg: str, *args) -> None:
msg (str): The warning message to log.
"""
_thread_id = threading.get_ident()
if not __verbose__:
logging.warning(f"{_thread_id}:{__version__} [W] {msg}", *args)
if not verbose:
logging.warning(f"{_thread_id} [W] {msg}", *args)
return
module, line_number, function_name = get_context_info()
if module and line_number and function_name:
logging.warning(
f"{_thread_id}:{__version__} [W] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
f"{_thread_id} [W] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
*args,
)
else:
logging.warning(f"{_thread_id}:{__version__} [W] {msg}", *args)
logging.warning(f"{_thread_id} [W] {msg}", *args)
def logging_info(msg: str, *args) -> None:
@@ -97,17 +94,17 @@ def logging_info(msg: str, *args) -> None:
msg (str): The info message to log.
"""
_thread_id = threading.get_ident()
if not __verbose__:
logging.info(f"{_thread_id}:{__version__} [*] {msg}", *args)
if not verbose:
logging.info(f"{_thread_id} [*] {msg}", *args)
return
module, line_number, function_name = get_context_info()
if module and line_number and function_name:
logging.info(
f"{_thread_id}:{__version__} [*] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
f"{_thread_id} [*] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
*args,
)
else:
logging.info(f"{_thread_id}:{__version__} [*] {msg}", *args)
logging.info(f"{_thread_id} [*] {msg}", *args)
def logging_debug(msg: str, *args) -> None:
@@ -120,9 +117,9 @@ 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.debug(
f"{_thread_id}:{__version__} [DBG] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
logging.info(
f"{_thread_id} [DBG] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
*args,
)
else:
logging.debug(f"{_thread_id}:{__version__} [DBG] {msg}", *args)
logging.debug(f"{_thread_id} [DBG] {msg}", *args)
+83 -93
View File
@@ -2,15 +2,12 @@
This module provides serialization and deserialization functions for Pydantic models that extend BaseModel
"""
import importlib
import json
import uuid
from datetime import date, datetime
from types import UnionType
from typing import Any, Dict, Union, get_args, get_origin
from urllib.parse import parse_qs, urlparse
import json
from datetime import datetime
from typing import Dict, Any
from pydantic import BaseModel
from urllib.parse import urlparse, parse_qs
from amqp.adapter.logging_utils import logging_debug
@@ -28,12 +25,8 @@ 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, _ 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)
for key, value in obj.model_dump().items():
if isinstance(value, uuid.UUID):
obj_dict[key] = str(value)
elif isinstance(value, datetime):
@@ -53,50 +46,8 @@ def serialize_object(obj: BaseModel) -> str:
else:
obj_dict[key] = value
dict_string = json.dumps(obj_dict)
return f"{module_name}.{class_name}:{dict_string}"
def python_type_to_json(data):
"""
Convert Python objects (dict, list, tuple, set, datetime, etc.) to a JSON string.
Handles common non-JSON-serializable types like datetime, date, and sets.
Args:
data: Python object to be converted (dict, list, tuple, set, etc.)
Returns:
str: JSON string representation of the input data
Raises:
TypeError: If the object contains non-serializable types (e.g., custom classes)
"""
def default_serializer(obj):
"""Custom serializer for non-JSON-serializable objects."""
if isinstance(obj, (datetime, date)):
return obj.isoformat() # Convert datetime/date to ISO string
elif isinstance(obj, set):
return list(obj) # Convert sets to lists
elif hasattr(obj, "__dict__"):
return obj.__dict__ # Attempt to serialize custom objects
else:
raise TypeError(f"Object of type {type(obj)} is not JSON serializable")
return json.dumps(data, default=default_serializer, indent=0, ensure_ascii=False)
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
dict_string = ",".join(f"{k}:{v}" for k, v in obj_dict.items())
return f"{class_name}:{{{dict_string}}}"
def deserialize_object(data: str) -> BaseModel:
@@ -111,54 +62,93 @@ def deserialize_object(data: str) -> BaseModel:
Returns:
An R2RSerializable object.
"""
class_path, dict_str = data.split(":", 1)
module_path, class_name = class_path.rsplit(".", 1)
class_name, dict_str = data.split(":", 1)
dict_str = dict_str.strip("{}")
pairs = dict_str.split(",") if dict_str else []
obj_dict: Dict[str, Any] = json.loads(dict_str)
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)
logging_debug(f"Deserialized object: {class_path}")
logging_debug(f"Deserialized object: {class_name}")
if "[" in class_name:
class_name = class_name.split("[")[0]
logging_debug(f"trying object: {class_name}")
module = importlib.import_module(module_path)
cls = getattr(module, class_name)
cls = globals().get(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)
+2 -9
View File
@@ -4,7 +4,7 @@ Helper functions used to aid with (de)serialization of messages sent over AMQP.
import struct
from io import BytesIO
from typing import Dict, List, Union
from typing import List, Union, Dict
def long_to_bytes(x: int) -> bytes:
@@ -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 struct.error:
except IOError:
return 0
@@ -51,10 +51,7 @@ def csv_as_list(input_str: str) -> List[str]:
def add_to_map(map_: Dict[str, List[str]], token: str) -> None:
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 :])
@@ -66,7 +63,3 @@ def parse_map_list_string(input_str: str) -> Dict[str, List[str]]:
for token in input_str[1:-1].split("],"):
add_to_map(map_, token)
return map_
def str_to_bool(value: str) -> bool:
return value.lower() in ("true", "yes", "1", "t", "y", "on")
+13 -4
View File
@@ -1,12 +1,13 @@
import base64
import logging
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.logging_utils import logging_error, logging_info
from amqp.adapter.serializer import map_as_string, parse_map_string
from amqp.adapter.trace_info_adapter import TraceInfoAdapter
from amqp.model.model import ServiceMessage
from amqp.model.service_message_type import ServiceMessageType
from amqp.model.snowflake_id import SnowflakeId
from amqp.adapter.serializer import parse_map_string, map_as_string
RS = "|"
SPLIT_RS = r"\|"
@@ -98,7 +99,9 @@ class ServiceMessageFactory:
i += 1
reply_to = input_str_array[i] if i < len(input_str_array) else ""
i += 1
trace_info = parse_map_string(input_str_array[i]) if i < len(input_str_array) else {}
trace_info = (
parse_map_string(input_str_array[i]) if i < len(input_str_array) else {}
)
i += 1
message = input_str_array[i] if i < len(input_str_array) else ""
return ServiceMessage(
@@ -108,7 +111,11 @@ class ServiceMessageFactory:
recipient_name=recipient_name,
reply_to=reply_to,
trace_info=trace_info,
message=(base64.b64decode(message).decode() if payload_type & 0x80 else message),
message=(
base64.b64decode(message).decode()
if payload_type & 0x80
else message
),
)
except Exception as e:
logging_error(
@@ -141,7 +148,9 @@ class ServiceMessageFactory:
:return: formatted message type
"""
num_type = (
message.message_type.value if message.message_type else ServiceMessageType.UNKNOWN.value
message.message_type.value
if message.message_type
else ServiceMessageType.UNKNOWN.value
) - 1
fmt = (num_type & 0x7F) | (message.payload_type & 0x80)
return f"{fmt:02X}"
-205
View File
@@ -1,205 +0,0 @@
""" """
import importlib
import json
from typing import Any
class TypeDescriptor:
def __init__(self, type_string: str, extra_init_data: dict):
self.type_string = type_string
self.extra_init_data = extra_init_data or {}
self.is_json: bool = False
self.type: str = self._extract_type()
def _extract_type(self) -> str:
if self.type_string.startswith("<class '"):
return self.type_string[8:-2]
# Handle Annotated case of FastAPI annotation
if "typing.Optional" in self.type_string:
annotated = _extract_type_of(self.type_string)
if "typing.Annotated" in annotated:
# Extract content between the first square brackets after Annotated
match = _extract_type_of(annotated)
if match:
# Split by comma and get first part, then strip whitespace
parts = [p.strip() for p in match.split(",")]
if len(parts) > 1:
self.is_json = parts[1].lower() == "json"
return parts[0] if parts else ""
else:
# Handle simple Optional case
return annotated
return _extract_type_of(self.type_string)
def instantiate(self, payload: Any, is_json: bool = False) -> Any:
return _instantiate(
(
self.type
if "typing.Optional" in self.type_string or self.type_string.startswith("<class ")
else self.type_string
),
payload,
self.is_json or is_json,
extra_init_data=self.extra_init_data,
)
def __repr__(self) -> str:
return f"TypeDescriptor(type='{self.type}', is_json={self.is_json})"
def _extract_type_of(input_string: str) -> str:
stack = []
start_index = -1
end_index = -1
for i, char in enumerate(input_string):
if i > 0: # this assumes paramerized type, that must be at lest 1 char long
if char == "[":
if start_index == -1:
start_index = i
stack.append(char)
elif char == "]":
if len(stack) == 1:
end_index = i
break
if len(stack) > 0:
stack.pop()
if start_index != -1 and end_index != -1:
return input_string[start_index + 1 : end_index].strip()
else:
return input_string.strip() if start_index == 1 and end_index == -1 else ""
def merge_extra_init_data(cls, json_payload, extra_init_data: dict):
"""
Merges extra_init_data into json_payload only when:
1. cls has an attribute named like the key
2. json_payload doesn't already have the key
"""
return {
**json_payload,
**{
k: v
for k, v in extra_init_data.items()
if (hasattr(cls, k) or (hasattr(cls, "model_fields") and k in cls.model_fields))
and k not in json_payload
},
}
def _instantiate(
type_string: str, payload: Any, is_json: bool = False, extra_init_data: dict = None
) -> Any:
"""
Dynamically imports and instantiates a class from a TypeDescriptor's type_string attribute.
Args:
type_string: TypeDescriptor instance with the class path
payload: Argument to pass to the constructor
Returns:
Instance of the target class initialized with payload
Raises:
ImportError: If the module/class cannot be imported
TypeError: If instantiation fails
"""
# Handle JSON payload first
json_payload = json.loads(payload) if is_json and isinstance(payload, str) else None
# Handle basic types
if type_string == "str":
return str(payload)
if type_string == "int":
return int(payload)
if type_string == "float":
return float(payload)
if type_string == "bool":
return bool(payload)
# Handle typing module types
if type_string.startswith("typing.") or type_string.startswith("list"):
if (
type_string.startswith("typing.List")
or type_string.startswith("typing.Sequence")
or type_string.startswith("list")
):
if json_payload is not None:
return json_payload
inner_type = _extract_type_of(type_string)
if inner_type:
return [
_instantiate(inner_type, p, extra_init_data=extra_init_data)
for p in (payload if isinstance(payload, list) else payload.split(","))
]
return payload if isinstance(payload, list) else payload.split(",")
if type_string.startswith("typing.Dict") or type_string.startswith("typing.Mapping"):
if json_payload is not None:
return json_payload
return dict(payload)
if type_string.startswith("typing.Set"):
if json_payload is not None:
return set(json_payload)
_set_type = _extract_type_of(type_string)
return set(
payload
if isinstance(payload, list)
else [_instantiate(_set_type, p) for p in payload.split(",")]
)
if type_string.startswith("typing.Tuple"):
if json_payload is not None:
return tuple(json_payload)
_tuple_type = _extract_type_of(type_string)
return tuple(
payload
if isinstance(payload, list)
else [_instantiate(_tuple_type, p) for p in payload.split(",")]
)
if type_string.startswith("typing.Optional"):
inner_type = _extract_type_of(type_string)
if inner_type:
return _instantiate(inner_type, payload, is_json, extra_init_data)
return payload
# Handle regular class instantiation
parts = type_string.split(".")
module_path = ".".join(parts[:-1])
class_name = parts[-1]
try:
# Import the module
module = importlib.import_module(module_path)
# Get the class
cls = getattr(module, class_name)
if isinstance(payload, cls):
return payload
# Instantiate with payload
if json_payload is not None:
_args = merge_extra_init_data(cls, json_payload, extra_init_data or {})
if hasattr(cls, "from_dict"):
return cls.from_dict(_args)
elif hasattr(cls, "create"):
return cls.create(**_args)
return cls(**_args)
if isinstance(payload, dict):
_args = merge_extra_init_data(cls, payload, extra_init_data or {})
if hasattr(cls, "from_dict"):
return cls.from_dict(_args)
elif hasattr(cls, "create"):
return cls.create(**_args)
return cls(**_args)
return cls(payload)
except ImportError as e:
raise ImportError(f"Could not import module '{module_path}'") from e
except AttributeError as e:
raise ImportError(f"Module '{module_path}' has no class '{class_name}'") from e
except TypeError as e:
raise TypeError(f"Failed to instantiate {type_string} with payload {payload}") from e
+12 -4
View File
@@ -41,7 +41,9 @@ class AMQAdapter:
fallback=False,
)
# these 2 shall be provided by Docker swarm, format is a random alphanumeric string
self.swarm_service_id = os.getenv("SWARM_SERVICE_ID", default="clever-amqp-service")
self.swarm_service_id = os.getenv(
"SWARM_SERVICE_ID", default="clever-amqp-service"
)
self.swarm_task_id = os.getenv("SWARM_TASK_ID", default="python-adapter")
def adapter_prefix(self) -> str:
@@ -66,9 +68,15 @@ class Dispatch:
:param config: config parser to read configuration values
"""
self.use_dlq = config.getboolean("CleverMicro-AMQ", self.PREFIX + "use-dlq", fallback=False)
self.amq_host = config.get("CleverMicro-AMQ", self.PREFIX + "amq-host", fallback="rabbitmq")
self.amq_port = config.getint("CleverMicro-AMQ", self.PREFIX + "amq-port", fallback=5672)
self.use_dlq = config.getboolean(
"CleverMicro-AMQ", self.PREFIX + "use-dlq", fallback=False
)
self.amq_host = config.get(
"CleverMicro-AMQ", self.PREFIX + "amq-host", fallback="rabbitmq"
)
self.amq_port = config.getint(
"CleverMicro-AMQ", self.PREFIX + "amq-port", fallback=5672
)
self.amq_port_tls = config.getint(
"CleverMicro-AMQ", self.PREFIX + "amq-port-tls", fallback=5671
)
+10 -21
View File
@@ -2,41 +2,30 @@
# CleverMicro AMQ Adapter settings
# ====================================================================
[CleverMicro-AMQ]
# CleverMicro AMQ Adapter settings that identify this AMQ Adapter instance to other adapters
# A name to identify this service for the AMQ Adapter
cm.amq-adapter.service-name=cleverthis-service-name
# The CleverMicro AMQ Adapter instance ID, used to identify this instance in the CleverMicro AMQ Adapter cluster
cm.amq-adapter.generator-id=1
# Route mapping. For route detail description see section 3.1 in Onboarding Guide for Python:
# https://docs.cleverthis.com/en/architecture/microservices/onboarding/python
# NOTE: this is a comma-separated list of route mappings (multiple, comma-separated routes are allowed here)
cm.amq-adapter.route-mapping=
# Indicate if AMQ Adapter should respond with HTTP 401 Unauthorized if the request does not contain valid JWT token
cm.amq-adapter.service-name=amq-adapter-python
cm.amq-adapter.generator-id=164
cm.amq-adapter.route-mapping=amq-cleverswarm::amq-cleverswarm:cleverswarm.localhost.#:5:0
cm.amq-adapter.require-authenticated-user=false
# CleverMicro AMQ Adapter dispatcher, settings related to RabbitMQ connection and optional REST forwarding
cm.dispatch.use-dlq=true
cm.dispatch.use-confirms=false
cm.dispatch.amq-host=rabbitmq
cm.dispatch.amq-port=5672
cm.dispatch.amq-port-tls=5671
cm.dispatch.download-dir=/tmp/downloaded_files
# Not used for CleveSwarm nor CleverBRAG or in tight coupling mode.
# Not used for CleverSwarm nor CleverBRAG or in tight coupling mode.
# applicable only in loose coupling mode, specify the destination where to forward the REST request
#cm.dispatch.service-host=cleverthis-service-name.localhost
#cm.dispatch.service-host=cleverswarm.localhost
#cm.dispatch.service-port=8080
# Backpressure monitor settings
#
# Define maximum number of messages of being processed in parallel before triggering backpressure OVERLOAD alert
cm.backpressure.threshold=5
# Define maximum CPU usage (%) before triggering backpressure OVERLOAD alert
cm.backpressure.threshold-cpu-overload=90
# Define maximum CPU usage (%) before triggering backpressure IDLE alert
cm.backpressure.threshold-cpu-idle=10
# Define maximum CPU usage before triggering backpressure OVERLOAD alert
cm.backpressure.threshold-cpu=90
# Define the time window between the Backpressure reports, in seconds
cm.backpressure.time-window=10
# Define the time window length, in seconds, during which, if CPU usage is consistently above the threshold, the system will trigger OVERLOAD alert
cm.backpressure.cpu-overload-duration=30
# Define the time window length, in seconds, during which, if no request is made to the Service, the system will trigger IDLE alert
cm.backpressure.cpu-idle-duration=30
cm.backpressure.idle-duration=30
+11
View File
@@ -0,0 +1,11 @@
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)
+69 -121
View File
@@ -1,10 +1,8 @@
import base64
import json
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List
from aio_pika.abc import AbstractRobustConnection
from typing import List, Dict
from dataclasses import dataclass, field
from amqp.model.service_message_type import ServiceMessageType
from amqp.model.snowflake_id import SnowflakeId
@@ -18,29 +16,9 @@ Define the structure of messages used in Clever Micro
"""
class DataMessageMagicByte(Enum):
HTTP_REQUEST = "A"
RPC_REQUEST = "R"
@dataclass
class CleverMicroMessage:
DEFAULT_CONTENT_TYPE = ["application/octet-stream"]
"""
The common structure of a message that transfers the user request to the Service and back. All data related messages
conform to this base structure:
magic - identifies the message type (REST, RPC, DataResponse)
id - unique message ID
m_field - method name (for REST style calls) or service name (for RPC style calls) or response code for DataResponse
d_field - domain name (for REST style calls) or module name (for RPC style calls) or error code for DataResponse
p_field - path (for REST style calls) or method name (for RPC style calls) or error cause for DataResponse
headers - headers of the message (e.g. Content-Type, Accept, etc.)
trace_info - trace information (e.g. trace ID, span ID, etc. for Jaeger trace monitor)
base64body - the body of the message encoded in base64
"""
def __init__(
self,
magic: str,
@@ -88,19 +66,11 @@ class CleverMicroMessage:
def body(self) -> bytes:
return base64.b64decode(self.base64body())
def content_type(self) -> str:
def contentType(self) -> str:
return self.headers().get("Content-Type", self.DEFAULT_CONTENT_TYPE)[0]
class DataMessage(CleverMicroMessage):
"""
The class represents a structure of a message that transfers the user request to the Service.
It may also be used to make a call/request between services.
The response to this call is in format of DataResponse class.
See DataMessageFactory for code to instantiate, serialize and deserialize the class
"""
def __init__(
self,
magic: str,
@@ -112,7 +82,9 @@ class DataMessage(CleverMicroMessage):
trace_info: Dict[str, str],
base64body: bytes,
):
super().__init__(magic, id, method, domain, path, headers, trace_info, base64body)
super().__init__(
magic, id, method, domain, path, headers, trace_info, base64body
)
def method(self) -> str:
return self.m_field()
@@ -123,77 +95,19 @@ class DataMessage(CleverMicroMessage):
def path(self) -> str:
return self.p_field()
def routing_key(self) -> str:
return sanitize_as_rabbitmq_name(f"{self.d_field()}.{self.p_field()}")
class MultipartDataMessage(DataMessage):
def __init__(self, message: DataMessage, extra_data: dict):
super().__init__(
message.magic(),
message.id(),
message.method(),
message.domain(),
message.path(),
message.headers(),
message.trace_info(),
message.base64body(),
)
self.extra_data = extra_data
@dataclass
class AMQMessage(DataMessage):
"""
Expand the core DataMessage type to include the channel, which is required to create a dedicated queue
in case the response should initiate file download.
"""
_connection: AbstractRobustConnection = field(repr=False) # Don't show in repr
def __init__(self, data_message: DataMessage, connection: AbstractRobustConnection):
# Initialize the parent DataMessage with all properties from the input data_message
super().__init__(
magic=data_message.magic(),
id=data_message.id(),
method=data_message.method(),
domain=data_message.domain(),
path=data_message.path(),
headers=data_message.headers(),
trace_info=data_message.trace_info(),
base64body=data_message.base64body(),
)
self._connection = connection
self.extra_data = data_message.extra_data if hasattr(data_message, "extra_data") else {}
@property
def connection(self) -> AbstractRobustConnection:
"""Get the AMQP connection associated with this message"""
return self._connection
class DataResponse(CleverMicroMessage):
"""
The response message to REST & RPC style calls (note: the request is made with DataMessage)
See DataResponseFactory for code to instantiate, serialize and deserialize the class
"""
def __init__(
self,
magic: str,
id: SnowflakeId,
response_code: str,
content_type: str,
error: str,
error_cause: str,
headers: Dict[str, List[str]],
trace_info: Dict[str, str],
base64body: bytes,
):
if headers is None:
headers = {}
headers["Content-Type"] = [content_type]
super().__init__(
magic,
id,
@@ -265,6 +179,65 @@ class AMQRoute:
return self.__str__()
@dataclass
class DataMessage:
"""
The class represents a structure of a message that transfers the user request to the Service.
It may also be used to make a call/request between services.
The response to this call is in format of DataResponse class.
Please see DataMessageFactory for code to instantiate, serialize and deserialize the class
"""
magic: str
id: SnowflakeId
method: str
domain: str
path: str
headers: dict[str, List[str]]
trace_info: dict[str, str]
base64_body: bytes
def routing_key(self) -> str:
return sanitize_as_rabbitmq_name(f"{self.domain}.{self.path}")
def body(self) -> bytes:
return base64.b64decode(self.base64_body)
class MultipartDataMessage(DataMessage):
def __init__(self, message: DataMessage, extra_data: dict):
super().__init__(
message.magic,
message.id,
message.method,
message.domain,
message.path,
message.headers,
message.trace_info,
message.base64_body,
)
self.extra_data = extra_data
@dataclass
class DataResponse:
"""
The response message to RPC style call (note: the request is made with DataMessage)
Please see DataResponseFactory for code to instantiate, serialize and deserialize the class
"""
id: str
response_code: int
content_type: str
response: bytes
error: str | None
error_cause: str | None
trace_info: dict[str, str]
def body(self) -> bytes:
return base64.b64decode(self.response)
# Ensure backward compatibility
AMQResponse = DataResponse
@@ -293,7 +266,7 @@ class ScalingRequestAlert(Enum):
IDLE = "IDLE"
UPDATE = "UPDATE"
SCALE_UP = "SCALE_UP"
SCALE_DOWN = "SCALE_DOWN"
SCSLE_DOWN = "SCALE_DOWN"
@dataclass(frozen=True)
@@ -328,32 +301,7 @@ class ScalingRequest:
task_id=data["taskId"],
max_availability=data["maxAvailability"],
current_availability=data["currentAvailability"],
request_type=ScalingRequestAlert(data["requestType"]), # Convert string to Enum
request_type=ScalingRequestAlert(
data["requestType"]
), # Convert string to Enum
)
class AMQErrorMessage:
"""
The error response message to RPC style call (note: the request is made with DataMessage)
"""
def __init__(
self,
response_code: int,
error: str,
detail: str,
content_type: str = "application/json",
):
self.response_code = response_code
self.error = error
self.detail = detail
self.content_type = content_type
def serialize(self) -> bytes:
return json.dumps(
{
"response_code": self.response_code,
"error": self.error,
"detail": self.detail,
}
).encode("utf-8")
+27 -18
View File
@@ -1,21 +1,20 @@
import logging
import aio_pika
from aio_pika.abc import AbstractRobustExchange, AbstractRobustQueue
from pika.exchange_type import ExchangeType
from aio_pika.abc import AbstractRobustQueue, AbstractRobustExchange, ExchangeType
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.logging_utils import (
logging_debug,
logging_error,
logging_debug,
logging_info,
logging_warning,
)
from amqp.model.model import AMQRoute, DataMessage
from amqp.model.model import DataMessage, AMQRoute
from amqp.router.router_base import (
AMQ_REPLY_TO_EXCHANGE,
AMQ_RPC_EXCHANGE,
DLQ_EXCHANGE,
AMQ_RPC_EXCHANGE,
)
from amqp.router.router_consumer import RouterConsumer
from amqp.router.utils import sanitize_as_rabbitmq_name
@@ -55,7 +54,7 @@ class RabbitMQClient:
name=cm_reply_to_queue, durable=True, exclusive=False, auto_delete=False
)
self.reply_to_exchange = await self.channel.declare_exchange(
AMQ_REPLY_TO_EXCHANGE, type=ExchangeType.direct
AMQ_REPLY_TO_EXCHANGE, type=ExchangeType.DIRECT
)
await self.reply_to_queue.bind(
exchange=self.reply_to_exchange, routing_key=cm_reply_to_queue
@@ -85,15 +84,21 @@ class RabbitMQClient:
After this the Service is listening on (consuming from) the queue for incoming messages.
"""
if self.router.is_open(self.channel):
queue_args = self.DLQ_ARGS if self.amq_configuration.dispatch.use_dlq else None
logging_info("RabbitMQClient register routes, cfg mapping: [%s]", route_mapping)
queue_args = (
self.DLQ_ARGS if self.amq_configuration.dispatch.use_dlq else None
)
logging_info(
"RabbitMQClient register routes, cfg mapping: [%s]", route_mapping
)
# convert the comma-separated string into list of AMQRoute objects
requested_routes = self.router.unwrap_route_list(route_mapping)
for route in requested_routes:
try:
# Note: AMQRoute allows one of the exchange or queue to be None, in which case the default is used
queue_name = (
route.queue if route.queue else self.router.get_consume_queue_name()
route.queue
if route.queue
else self.router.get_consume_queue_name()
)
logging_info("RabbitMQClient register route: [%s]", route)
queue: AbstractRobustQueue = await self.channel.declare_queue(
@@ -103,9 +108,13 @@ class RabbitMQClient:
auto_delete=False,
arguments=queue_args,
)
_exchange_name: str = route.exchange if route.exchange else AMQ_RPC_EXCHANGE
exchange: AbstractRobustExchange = await self.channel.declare_exchange(
name=_exchange_name, type=ExchangeType.topic
_exchange_name: str = (
route.exchange if route.exchange else AMQ_RPC_EXCHANGE
)
exchange: AbstractRobustExchange = (
await self.channel.declare_exchange(
name=_exchange_name, type=ExchangeType.TOPIC
)
)
await queue.bind(exchange, route.key)
_c_tag = await queue.consume(callback, no_ack=False)
@@ -190,14 +199,16 @@ class RabbitMQClient:
message
), # context path is a routing key
)
logging_info(f"Msg Publish (no-reply), messageId: {message.id()}, to: {route}")
logging_info(
f"Msg Publish (no-reply), messageId: {message.id}, to: {route}"
)
else:
# This is RPC call, there should be response
pika_message: aio_pika.message.Message = aio_pika.message.Message(
body=DataMessageFactory.serialize(message),
content_encoding="utf-8",
delivery_mode=2,
correlation_id=str(message.id()),
correlation_id=str(message.id),
reply_to=self.router.get_reply_to_queue_name(),
)
await self.rpc_exchange.publish(
@@ -208,14 +219,12 @@ class RabbitMQClient:
)
logging_info(
"Msg Publish (RPC call), messageId: %s, to: %s",
message.id().as_string(),
message.id.as_string(),
route.as_string(),
)
else:
logging_warning(
"Message not sent, no route for %s/%s",
message.domain(),
message.path(),
"Message not sent, no route for %s/%s", message.domain, message.path
)
except Exception as err:
+7 -2
View File
@@ -1,3 +1,4 @@
import logging
import re
from typing import Set
@@ -40,9 +41,13 @@ class RouteDatabase:
return re.compile(".*")
return re.compile(regex)
def matches(self, routing_key_from_route: str, routing_key_from_message: str) -> bool:
def matches(
self, routing_key_from_route: str, routing_key_from_message: str
) -> bool:
route_pattern = self.pattern(routing_key_from_route)
return bool(route_pattern.match(sanitize_as_rabbitmq_name(routing_key_from_message)))
return bool(
route_pattern.match(sanitize_as_rabbitmq_name(routing_key_from_message))
)
def wrap_routes(self) -> str:
return ",".join(str(route) for route in self.routes)
+19 -8
View File
@@ -1,7 +1,8 @@
import logging
from typing import Callable, List, Set
from amqp.adapter.amq_route_factory import NULL_ROUTE, AMQRouteFactory
from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
from amqp.adapter.amq_route_factory import AMQRouteFactory, NULL_ROUTE
from amqp.adapter.logging_utils import logging_error, logging_debug, logging_info
from amqp.model.model import AMQRoute, DataMessage
from amqp.router.route_database import RouteDatabase
from amqp.router.utils import sanitize_as_rabbitmq_name
@@ -27,7 +28,7 @@ class RouterBase:
self.route_removed_notifier: Callable[[AMQRoute], None] | None = None
def get_routing_key(self, message: DataMessage) -> str:
return sanitize_as_rabbitmq_name(f"{message.domain()}.{message.path()}")
return sanitize_as_rabbitmq_name(f"{message.domain}.{message.path}")
def wrap_routes(self) -> str:
return self.route_database.wrap_routes()
@@ -37,7 +38,10 @@ class RouterBase:
def unwrap_route_list(self, route_string: str) -> List[AMQRoute]:
logging_debug("RouterMaster.unwrapRoutes, route(s): %s", route_string)
routes = [AMQRouteFactory.from_string(route_str) for route_str in route_string.split(",")]
routes = [
AMQRouteFactory.from_string(route_str)
for route_str in route_string.split(",")
]
return [route for route in routes if route != NULL_ROUTE]
def get_routes(self) -> Set[AMQRoute]:
@@ -47,12 +51,14 @@ class RouterBase:
return self.route_database.find_route(domain, path)
def find_route_by_message(self, message: DataMessage) -> AMQRoute | None:
return self.route_database.find_route(message.domain(), message.path())
return self.route_database.find_route(message.domain, message.path)
def add_routes(self, message: str) -> None:
try:
amq_route_list = self.unwrap_route_list(message)
logging_info(" [DBG] RouterMaster.addRoute [%d]: %s", len(amq_route_list), message)
logging_info(
" [DBG] RouterMaster.addRoute [%d]: %s", len(amq_route_list), message
)
for route in amq_route_list:
logging_info("RouterMaster.addRoute, adding route: %s", route)
if not route.exchange and not route.queue and not route.key:
@@ -76,9 +82,14 @@ class RouterBase:
def remove_routes(self, message: str) -> None:
try:
to_remove = [AMQRouteFactory.from_string(route_str) for route_str in message.split(",")]
to_remove = [
AMQRouteFactory.from_string(route_str)
for route_str in message.split(",")
]
removed_count = sum(1 for route in to_remove if self.remove_route(route))
logging_info(f"RouterMaster.removeRoutes(): {message}, removed={removed_count}")
logging_info(
f"RouterMaster.removeRoutes(): {message}, removed={removed_count}"
)
except IOError as err:
logging_error(f"Can't remove route(s): {err}")
+29 -13
View File
@@ -1,26 +1,26 @@
import logging
from typing import Set
from aio_pika.abc import (
AbstractIncomingMessage,
AbstractRobustChannel,
AbstractRobustQueue,
AbstractIncomingMessage,
)
from amqp.adapter.amq_route_factory import AMQRouteFactory
from amqp.adapter.logging_utils import (
logging_debug,
logging_error,
logging_info,
logging_debug,
logging_warning,
)
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQRoute, ServiceMessage
from amqp.model.model import ServiceMessage, AMQRoute
from amqp.model.service_message_type import ServiceMessageType
from amqp.router.router_base import (
ANY_RECIPIENT,
CM_C_REQ_QUEUE,
CM_REQUEST_EXCHANGE,
CM_RESPONSE_EXCHANGE,
ANY_RECIPIENT,
)
from amqp.router.router_producer import RouterProducer
@@ -77,13 +77,17 @@ class RouterConsumer(RouterProducer):
except Exception as ioe:
logging_error("RouterConsumer FATAL: cannot initialize, %s", ioe)
else:
logging_error("RouterConsumer FATAL: cannot initialize, RabbitMQ channel is not open.")
logging_error(
"RouterConsumer FATAL: cannot initialize, RabbitMQ channel is not open."
)
"""
* ServiceMessageType.ROUTING_DATA_REQUEST handler listening on queue 'CleverMicroRequest'
"""
async def handle_routing_data_request_message(self, message: AbstractIncomingMessage):
async def handle_routing_data_request_message(
self, message: AbstractIncomingMessage
):
# some debug statements, TODO - remove for prod release
delivery_tag = message.delivery_tag
debug = "Delivery.Properties = " + message.properties.__str__()
@@ -97,7 +101,9 @@ class RouterConsumer(RouterProducer):
await message.ack(False)
cm_message: ServiceMessage = self.service_message_factory.from_bytes(message.body)
cm_message: ServiceMessage = self.service_message_factory.from_bytes(
message.body
)
route_mapping: str = self.wrap_own_routes()
logging_debug(
"******[%s] (ROUTING_DATA_REQ) msg=%s, routeMap=[%s]",
@@ -118,7 +124,9 @@ class RouterConsumer(RouterProducer):
route_mapping,
cm_message.reply_to,
)
await self.publish_service_message(service_message, self.cm_response_exchange)
await self.publish_service_message(
service_message, self.cm_response_exchange
)
except Exception as err:
logging_error(
"******[%s] (ROUTING_DATA_REQ) Response to queue: %s failed: %s",
@@ -150,7 +158,9 @@ class RouterConsumer(RouterProducer):
serviceMessage: ServiceMessage = self.service_message_factory.of(
ServiceMessageType.ROUTING_DATA, route.as_string(), ANY_RECIPIENT
)
await self.publish_service_message(serviceMessage, self.cm_response_exchange)
await self.publish_service_message(
serviceMessage, self.cm_response_exchange
)
self.add_own_route(route)
logging_info("RouterConsumer registered Route: %s", route)
else:
@@ -172,7 +182,9 @@ class RouterConsumer(RouterProducer):
service_message: ServiceMessage = self.service_message_factory.of(
ServiceMessageType.ROUTING_DATA_REMOVE, route.as_string(), ANY_RECIPIENT
)
if not await self.publish_service_message(service_message, self.cm_response_exchange):
if not await self.publish_service_message(
service_message, self.cm_response_exchange
):
logging_warning(
"RouterConsumer couldn't remove current route: %s, channel not open.",
route,
@@ -189,9 +201,13 @@ class RouterConsumer(RouterProducer):
* Cleanup - de-register routes with master.
"""
routes: Set[AMQRoute] = self.get_routes()
logging_info("RouterConsumer.cleanup() [%d]: %s", len(routes), self.wrap_routes())
logging_info(
"RouterConsumer.cleanup() [%d]: %s", len(routes), self.wrap_routes()
)
not_removed = sum(
1 for r in (self.remove_route_on_cleanup(route) for route in routes) if not r
1
for r in (self.remove_route_on_cleanup(route) for route in routes)
if not r
)
if not_removed > 0:
logging_warning(
+31 -19
View File
@@ -7,20 +7,21 @@
* forwards it to the application.
"""
import logging
from aio_pika import Message
from aio_pika.abc import (
AbstractIncomingMessage,
AbstractRobustChannel,
AbstractRobustQueue,
AbstractRobustConnection,
AbstractRobustExchange,
AbstractRobustQueue,
AbstractIncomingMessage, ExchangeType,
)
from pika.exchange_type import ExchangeType
from amqp.adapter.logging_utils import (
logging_debug,
logging_error,
logging_info,
logging_debug,
logging_warning,
)
from amqp.adapter.service_message_factory import ServiceMessageFactory
@@ -28,13 +29,13 @@ from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import ServiceMessage
from amqp.model.service_message_type import ServiceMessageType
from amqp.router.router_base import (
ANY_RECIPIENT,
CM_C_AMQ_QUEUE,
CM_P_REPLY_TO,
CM_P_RESP_QUEUE,
RouterBase,
CM_REQUEST_EXCHANGE,
CM_RESPONSE_EXCHANGE,
RouterBase,
CM_C_AMQ_QUEUE,
CM_P_RESP_QUEUE,
CM_P_REPLY_TO,
ANY_RECIPIENT,
)
@@ -47,14 +48,16 @@ class RouterProducer(RouterBase):
"""
super().__init__()
self.connection: AbstractRobustConnection | None = None
self.service_message_factory = ServiceMessageFactory(configuration.amq_adapter.service_name)
self.service_message_factory = ServiceMessageFactory(
configuration.amq_adapter.service_name
)
self.amq_configuration: AMQConfiguration = configuration
self.channel: AbstractRobustChannel | None = None
self.cm_request_exchange: AbstractRobustExchange | None = None
self.cm_response_exchange: AbstractRobustExchange | None = None
logging_info(
"RouterProducer.<init>: %s, route: %s",
f"RouterProducer.<init>: %s, route: %s",
self.amq_configuration.amq_adapter.service_name,
self.amq_configuration.amq_adapter.route_mapping,
)
@@ -78,13 +81,13 @@ class RouterProducer(RouterBase):
# **** set up the service discovery internal exchanges ****
# 1. declare RequestExchange to where to publish the RoutingData requests
self.cm_request_exchange = await self.channel.declare_exchange(
name=CM_REQUEST_EXCHANGE, type=ExchangeType.fanout
name=CM_REQUEST_EXCHANGE, type=ExchangeType.FANOUT
)
#
# 2. declare ResponseExchange where to publish current RoutingData (response to
# a RoutingData request)
self.cm_response_exchange = await self.channel.declare_exchange(
name=CM_RESPONSE_EXCHANGE, type=ExchangeType.fanout
name=CM_RESPONSE_EXCHANGE, type=ExchangeType.FANOUT
)
# 2a. declare a Response queue for Routing Data replies from Services
_cm_response_queue: str = self.get_response_queue_name()
@@ -96,7 +99,9 @@ class RouterProducer(RouterBase):
arguments={},
)
# 2b. bind queue to CM_RESPONSE_EXCHANGE w/o any conditions (empty routing key)
await _response_queue.bind(exchange=CM_RESPONSE_EXCHANGE, routing_key="#")
await _response_queue.bind(
exchange=CM_RESPONSE_EXCHANGE, routing_key="#"
)
# 2c. listen for incoming RoutingData messages
_c_tag = await _response_queue.consume(
callback=self.handle_routing_data_message, no_ack=False
@@ -110,7 +115,9 @@ class RouterProducer(RouterBase):
logging_error("RouterProducer FATAL: cannot initialize, %s", ioe)
else:
logging_error("RouterProducer FATAL: cannot initialize, RabbitMQ channel is not open.")
logging_error(
"RouterProducer FATAL: cannot initialize, RabbitMQ channel is not open."
)
async def handle_routing_data_message(self, message: AbstractIncomingMessage):
"""
@@ -118,13 +125,15 @@ class RouterProducer(RouterBase):
* routing data, invokes addRoute or removeRoute depending on the type of the message.
"""
logging_info(
"[%s] Received ROUTING DATA: [%s], msgId=%s",
f"[%s] Received ROUTING DATA: [%s], msgId=%s",
message.delivery_tag,
message.body,
message.message_id,
)
await message.ack(multiple=False)
cm_message: ServiceMessage = self.service_message_factory.from_bytes(message.body)
cm_message: ServiceMessage = self.service_message_factory.from_bytes(
message.body
)
# ignore the message if it comes from ourselves (even if from different instance)
logging_debug(
"******[%s] ROUTING_DATA, msg=%s, routeMap=[%s] sn=%s, replyTo=%s",
@@ -136,7 +145,8 @@ class RouterProducer(RouterBase):
)
if (
cm_message.is_valid()
and not self.amq_configuration.amq_adapter.service_name == cm_message.reply_to
and not self.amq_configuration.amq_adapter.service_name
== cm_message.reply_to
):
if cm_message.message_type == ServiceMessageType.ROUTING_DATA:
self.add_routes(cm_message.message)
@@ -155,7 +165,9 @@ class RouterProducer(RouterBase):
serviceMessage: ServiceMessage = self.service_message_factory.of(
ServiceMessageType.ROUTING_DATA_REQUEST, "", ANY_RECIPIENT
)
if not await self.publish_service_message(serviceMessage, self.cm_request_exchange):
if not await self.publish_service_message(
serviceMessage, self.cm_request_exchange
):
logging_warning(
"RouterProducer could not request route data: %s channel is not open.",
CM_REQUEST_EXCHANGE,
+40 -32
View File
@@ -1,32 +1,32 @@
import asyncio
import os
import threading
import time
from asyncio import AbstractEventLoop, Future
import os
from asyncio import Future, AbstractEventLoop
from typing import Dict
from aio_pika import Message
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustExchange
from opentelemetry import trace
from opentelemetry.trace import Tracer, TraceState
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from opentelemetry.trace import Tracer, TraceState
from amqp.adapter.backpressure_handler import BackpressureHandler
from amqp.adapter.cleverthis_service_adapter import AMQMessage, CleverThisServiceAdapter
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.data_response_factory import DataResponseFactory
from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
from amqp.adapter.serializer import map_as_string
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter, AMQMessage
from amqp.adapter.file_handler import StreamingFileHandler, FileHandler
from amqp.adapter.logging_utils import logging_error, logging_debug, logging_info
from amqp.adapter.service_message_factory import ServiceMessageFactory
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import (
DataMessage,
DataMessageMagicByte,
DataResponse,
DataMessage,
MultipartDataMessage,
)
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient
from amqp.adapter.serializer import map_as_string
def collect_trace_info():
@@ -75,15 +75,12 @@ class DataMessageHandler:
self.file_handler: FileHandler = file_handler
self.backpressure_handler: BackpressureHandler = backpressure_handler
async def inbound_data_message_callback_as_thread(self, message: AbstractIncomingMessage):
"""
The callback method that is called when a DataMessage is received from RabbitMQ, starts new Thread.
"""
async def inbound_data_message_callback(self, message: AbstractIncomingMessage):
threading.Thread(
target=lambda: asyncio.run(self.inbound_data_message_callback_no_thread(message))
target=lambda: asyncio.run(self.run_callback_thread(message))
).start()
async def inbound_data_message_callback_no_thread(self, message: AbstractIncomingMessage):
async def run_callback_thread(self, message: AbstractIncomingMessage):
"""
Receives the DataMessage as bytes from RabbitMQ queue and deserializes it to the DataMessage object
Ensures Jaeger trace-id propagation.
@@ -95,17 +92,19 @@ class DataMessageHandler:
amq_message = await self.reconstitute_data_message(message)
if amq_message is not None:
self.ensure_trace_info(amq_message)
if amq_message.magic() == DataMessageMagicByte.HTTP_REQUEST.value:
if amq_message.magic == DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST:
await self.run_http_request_magic(message, amq_message, self.loop)
else:
logging_error(
"Unknown MAGIC value: [%s], don't know how to interpret the message",
amq_message.magic(),
amq_message.magic,
)
asyncio.run_coroutine_threadsafe(message.nack(requeue=False), self.loop)
return
else:
logging_error("######[%s] No valid AMQ Message present.", message.delivery_tag)
logging_error(
"######[%s] No valid AMQ Message present.", message.delivery_tag
)
self.record_duration(start, message.delivery_tag)
asyncio.run_coroutine_threadsafe(message.ack(multiple=False), self.loop)
@@ -169,7 +168,9 @@ class DataMessageHandler:
amq_message = await DataMessageFactory.from_stream(
message.body, rabbitmq_url=self.rabbit_mq_url, loop=self.loop
)
logging_info(f"######[{delivery_tag}] AMQ Message from stream: {amq_message}")
logging_info(
f"######[{delivery_tag}] AMQ Message from stream: {amq_message}"
)
if amq_message is not None:
extra_data = await self.file_handler.on_message(
message=message,
@@ -178,24 +179,26 @@ class DataMessageHandler:
loop=self.loop,
output_dir=self.output_dir,
)
logging_info(f"######[{delivery_tag}] Multipart download result: {extra_data}")
logging_info(
f"######[{delivery_tag}] Multipart download result: {extra_data}"
)
amq_message = MultipartDataMessage(amq_message, extra_data)
if amq_message:
logging_info(
"######[%s] AMQ Message received, ctag=%s, msgId=%s, %s:%s, traceInfo=%s",
message.delivery_tag,
message.consumer_tag,
amq_message.id(),
amq_message.method(),
amq_message.path(),
map_as_string(amq_message.trace_info()),
amq_message.id,
amq_message.method,
amq_message.path,
map_as_string(amq_message.trace_info),
)
return amq_message
def ensure_trace_info(self, amq_message: DataMessage):
propagator = TraceContextTextMapPropagator()
context = propagator.extract(
carrier=amq_message.trace_info(), context=trace.context_api.get_current()
carrier=amq_message.trace_info, context=trace.context_api.get_current()
)
self.tracer.start_span(name="AMQ-Adapter-Consumer", context=context)
# logging_info("CTX=%s", context)
@@ -204,7 +207,7 @@ class DataMessageHandler:
# logging.info(" [*] CTX2=%s", context_with_span)
# propagator.inject(context_with_span, amq_message.trace_info, self.set_text)
# logging.info(" [***] NEW TI=%s CTX=%s SCOPE=%s",
# map_as_string(amq_message.trace_info()), context_with_span, context)
# map_as_string(amq_message.trace_info), context_with_span, context)
def record_duration(self, start, delivery_tag):
global last_data_message_time
@@ -226,15 +229,17 @@ class DataMessageHandler:
if reply_to:
pika_message: Message = Message(
body=DataResponseFactory.serialize(response),
correlation_id=response.id(),
correlation_id=response.id,
content_type=(
response.content_type()[0]
if isinstance(response.content_type(), list)
else response.content_type()
response.content_type[0]
if isinstance(response.content_type, list)
else response.content_type
),
)
_res = asyncio.run_coroutine_threadsafe(
self.reply_to_exchange.publish(message=pika_message, routing_key=reply_to),
self.reply_to_exchange.publish(
message=pika_message, routing_key=reply_to
),
self.loop,
)
logging_debug(
@@ -252,12 +257,15 @@ class DataMessageHandler:
:return:
"""
reply_id = message.correlation_id
content_type = message.content_type
delivery_tag = message.delivery_tag
debug = f"######[{delivery_tag}] (RPC REPLY) env: {message.headers}, props: {message.properties}, BODY: {message.body.decode()}"
logging_debug(debug)
# TODO - figure out implementation of this part, as this is sending reply back either to a service,
# or to the user
logging_info(f"OUTSTANDING RECV: {self.outstanding}, size={len(self.outstanding)}")
logging_info(
f"OUTSTANDING RECV: {self.outstanding}, size={len(self.outstanding)}"
)
future: Future = self.outstanding.pop(reply_id, None)
# if reply is None:
+11 -12
View File
@@ -1,23 +1,25 @@
import asyncio
import logging
import logging.config
import os
import asyncio
from asyncio import Future
from threading import Thread
from typing import Optional
from aio_pika.abc import AbstractRobustExchange
from amqp.adapter.amq_route_factory import AMQRouteFactory
from amqp.adapter.backpressure_handler import BackpressureHandler
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.amq_route_factory import AMQRouteFactory
from amqp.adapter.logging_utils import logging_info, logging_warning
from amqp.adapter.service_message_factory import ServiceMessageFactory
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 TracingConfig, initialize_telemetry
from amqp.service.tracing import initialize_trace
logging.basicConfig(level=logging.DEBUG)
logging.info(f"CWD = {os.getcwd()}")
@@ -46,12 +48,7 @@ class AMQService:
self.amq_configuration.amq_adapter.generator_id
)
self.rabbit_mq_client = RabbitMQClient(self.amq_configuration)
self.tracer = initialize_telemetry(
tracing_config=TracingConfig(
service_name=amq_configuration.amq_adapter.service_name,
service_version=os.getenv("SERVICE_VERSION", "unknown"),
)
)
self.tracer = initialize_trace(name=amq_configuration.amq_adapter.service_name)
self.amq_data_message_handler = None
self.service_message_factory = ServiceMessageFactory(
self.amq_configuration.amq_adapter.service_name
@@ -69,7 +66,9 @@ class AMQService:
self.loop.run_forever()
async def init(self, loop, service_adapter):
_reply_to_exchange: AbstractRobustExchange = await self.rabbit_mq_client.init(loop)
_reply_to_exchange: AbstractRobustExchange = await self.rabbit_mq_client.init(
loop
)
backpressure_handler = BackpressureHandler(
channel=self.rabbit_mq_client.channel,
loop=loop,
@@ -108,7 +107,7 @@ class AMQService:
try:
await self.rabbit_mq_client.register_inbound_routes(
route_mapping,
self.amq_data_message_handler.inbound_data_message_callback_no_thread,
self.amq_data_message_handler.inbound_data_message_callback,
)
return await self.rabbit_mq_client.register_data_reply_callback(
self.amq_data_message_handler.reply_received_callback
@@ -123,7 +122,7 @@ class AMQService:
return None
def send_message(self, message: DataMessage, future: Future):
self.amq_data_message_handler.outstanding[str(message.id())] = future
self.amq_data_message_handler.outstanding[str(message.id)] = future
def remove_outstanding_future(self, message_id: str):
self.amq_data_message_handler.outstanding.pop(message_id, None)
+88 -182
View File
@@ -1,196 +1,102 @@
"""
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.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 opentelemetry.sdk.trace import TracerProvider, Tracer
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from prometheus_client import start_http_server
from opentelemetry.exporter.prometheus import PrometheusMetricReader
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
logger = logging.getLogger(__name__)
import os
@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."""
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: config.service_name,
SERVICE_VERSION: config.service_version,
SERVICE_NAME: (
name if name else os.environ.get("APPLICATION_NAME", "Python App")
)
}
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."""
)
#
# 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.
provider = TracerProvider(resource=resource)
# Add OTLP exporter for Jaeger
otlp_exporter = setup_otlp_exporter(config)
provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
# Add console exporter in debug mode
if config.debug:
provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))
return provider
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)
# 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)
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")
# logging.info(' * Instrumenting application..... ')
# if Flask, then use:
# FlaskInstrumentor().instrument_app(flask_app)
# if FastAPI, then use
# FastAPIInstrumentor.instrument_app(app)
# Return configured tracer
return trace.get_tracer(
tracing_config.service_name,
tracing_config.service_version,
)
#
# 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.
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")
# 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)
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 {},
)
return trace.get_tracer(name)
-125
View File
@@ -1,125 +0,0 @@
"""
This module provides the CleverBRAG AMQP Adapter for handling API requests via AMQP.
The code IS MEANT to be part of CleverBRAG codebase, and has direct dependencies on the CleverBRAG codebase.
It is NOT MEANT to be used as a standalone module. It WILL NOT run outside the CleverBRAG context.
"""
from threading import Thread
from typing import Any, List
from uuid import UUID
# CleverBrag specific imports
import core
from fastapi.routing import APIRoute
from fastapi.security.http import HTTPAuthorizationCredentials
from pydantic import BaseModel
from shared.abstractions import R2RSerializable, User
# AMQP specific imports
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter
from amqp.config.amq_configuration import AMQConfiguration
from amqp.service.amq_service import AMQService
class BRAGArgument:
"""
This class is used to represent a BRAG API input argument. It is uses FastAPI annotation to define the arguments
(the argument type) that are passed as input arguments to the BRAG endpoint.
"""
def __init__(self, arg: Any):
self.arg = arg
self.module = arg.__module__
self.cls = arg.__name__ if hasattr(arg, "__name__") else str(arg)
self.is_r2r_serializable = issubclass(arg, R2RSerializable)
self.is_builtin = self.module == "builtins"
self.is_uuid = self.cls == "UUID"
self.is_optional = self.cls == "Optional"
self.is_list = self.cls == "list"
self.is_dict = self.cls == "dict"
self.is_base_model = issubclass(arg, BaseModel)
self.is_nested = self.is_list or self.is_optional or self.is_dict
if self.is_nested and hasattr(arg, "__args__"):
self.nested = BRAGArgument(getattr(arg, "__args__")[0])
if self.is_dict:
self.nested = (self.nested, BRAGArgument(getattr(arg, "__args__")[1]))
else:
self.nested = None
self.is_nested = False
def __repr__(self):
return f"{self.module}.{self.cls}" + (f"[{self.nested}]" if self.is_nested else "")
def convert_to_argument(arg: BRAGArgument, value: str) -> Any:
"""
Converts the argument to a string representation.
"""
try:
if arg.cls == "str":
return value
elif arg.cls == "bool":
return str(value).lower() == "true"
elif arg.cls == "float":
return float(value)
elif arg.cls == "int":
return int(value)
elif arg.is_uuid:
return UUID(value)
elif arg.is_optional:
return convert_to_argument(arg.nested, value) if arg.nested else None
elif arg.is_list:
return (
[convert_to_argument(arg.nested, v.strip()) for v in value.split(",")]
if arg.nested
else value.split()
)
elif arg.is_dict:
return (
{
convert_to_argument(arg.nested[0], k.strip()): convert_to_argument(
arg.nested[1], v.strip()
)
for k, v in (item.split(":") for item in value.split(","))
}
if arg.nested
else dict(item.split(":") for item in value.split(","))
)
elif arg.is_base_model:
return arg.arg.parse_raw(value)
else:
print(f"Unsupported argument type: {type(arg)}")
except Exception as e:
print(f"Error converting argument {arg}: {e}")
return None
class CleverBragAmqpAdapter(CleverThisServiceAdapter):
"""
CleverBRAG AMQP Adapter for CleverBRAG API. Implements the interface to CleverBRAG API.
"""
def __init__(
self,
amq_configuration: AMQConfiguration,
routes: List[APIRoute],
auth_provider: core.base.providers.auth.AuthProvider,
extra_init_data: dict,
):
super().__init__(
amq_configuration, routes, None, "auth_user", extra_init_data=extra_init_data
)
self.auth_provider = auth_provider
_amq_service: AMQService = AMQService(amq_configuration, self)
self.thread = Thread(target=_amq_service.run)
self.thread.start()
async def get_current_user(self, token: str) -> User:
"""
Overrides the default get_current_user method to use the token from the AMQP message.
:param token: JWT token from the AMQP message
:return: User object
"""
return await self.auth_provider.auth_wrapper(
HTTPAuthorizationCredentials(scheme="Bearer", credentials=token)
)
-44
View File
@@ -1,44 +0,0 @@
from clever_brag_adapter import BRAGArgument, CleverBragAmqpAdapter
from amqp.config.amq_configuration import AMQConfiguration
class R2RApp:
pass
def create_r2r_app():
return R2RApp()
if __name__ == "__main__":
#
# r2rapp: R2RApp = await create_r2r_app()
#
# In the standalone test context, need to start the loop as well. In the actual BRAG context,
# use 'await' instead.
r2rapp: R2RApp = create_r2r_app()
adapter: CleverBragAmqpAdapter = CleverBragAmqpAdapter(
AMQConfiguration("./application.properties.local"),
r2rapp.app.routes,
r2rapp.providers.auth,
{"app": r2rapp.config.app},
)
for method, routes in adapter.routes.items():
print(f"Method: {method}")
for _route in routes:
print(
f" Route: {_route.path} -> {_route.endpoint.__module__}.{_route.endpoint.__name__}"
)
_args = getattr(_route.endpoint, "__annotations__", {})
_verified_args = []
for arg in _args:
_arg = BRAGArgument(_args[arg])
_verified_args.append(_arg)
print(f" Arguments: {', '.join(map(str, _verified_args))}")
print("---------------------------")
print("Press ^C to exit...")
adapter.thread.join()
-149
View File
@@ -1,149 +0,0 @@
echo "#!/bin/bash"
. ./test_shared_functions.sh
extract_document_id() {
R=$1
C=$(echo $R | sed 's/[0-9][0-9][0-9]$//')
#echo ">>${C}<<"
if [[ ! -z "$C" ]] && [[ "$C" == *results* ]]; then
COLLECTION_IDS=($(echo $C | jq -r '.results[] | select(.document_count >= 1) | .id' | tr '\n' ' '))
echo "CIDs = ${COLLECTION_IDS[*]}"
else
echo "No IDs detected in the response."
fi
HTTP_RETURN_CODE=$(echo echo $1 | tail -1 | sed -e 's/^.*\([0-9][0-9][0-9]\)$/\1/')
}
echo "# 1. Create Collection Tests"
echo "=== Testing Collection Creation ==="
echo "# Test 1.1: Create collection with name only"
make_request "POST" "/collections" \
"-d '{\"name\":\"Test Collection\"}'"
echo "# Test 1.2: Create collection with name and description"
make_request "POST" "/collections" \
"-d '{\"name\":\"Test Collection\",\"description\":\"A test collection\"}'"
echo "# Test 1.3: Create collection with empty name (should fail)"
make_request "POST" "/collections" \
"-d '{\"name\":\"\"}'"
echo "# Test 1.4: Create collection with very long name (edge case)"
make_request "POST" "/collections" \
"-d '{\"name\":\"$(printf 'a%.0s' {1..1000})\"}'"
echo "# 2. List Collections Tests"
echo "=== Testing Collection Listing ==="
echo "# Test 2.1: List all collections"
make_request "GET" "/collections" ""
ALL_COLLECTION_IDS=("${COLLECTION_IDS[@]}")
echo "# Test 2.2: List collections with pagination"
make_request "GET" "/collections?offset=0&limit=10" ""
echo "# Test 2.3: List collections with invalid pagination (edge cases)"
make_request "GET" "/collections?offset=-1&limit=10" ""
make_request "GET" "/collections?offset=0&limit=0" ""
make_request "GET" "/collections?offset=0&limit=1001" ""
echo "# Test 2.4: List specific collections by IDs"
# Join array elements with commas
COMMA_SEPARATED_IDS=$(echo "${ALL_COLLECTION_IDS[*]}" | sed -e "s/ /,/")
make_request "GET" "/collections?ids=${COMMA_SEPARATED_IDS}" ""
echo "# 3. Get Collection Tests"
echo "=== Testing Collection Retrieval ==="
echo "# Test 3.1: Get collection by ID"
make_request "GET" "/collections/${ALL_COLLECTION_IDS[0]}" ""
echo "# Test 3.2: Get non-existent collection"
make_request "GET" "/collections/00000000-0000-0000-0000-000000000000" ""
echo "# Test 3.3: Get collection with invalid UUID format"
make_request "GET" "/collections/invalid-uuid" ""
echo "# 4. Update Collection Tests"
echo "=== Testing Collection Updates ==="
echo "# Test 4.1: Update collection name"
make_request "POST" "/collections/${ALL_COLLECTION_IDS[0]}" \
"-d '{\"name\":\"Updated Collection Name\"}'"
echo "# Test 4.2: Update collection description"
make_request "POST" "/collections/${ALL_COLLECTION_IDS[0]}" \
"-d '{\"description\":\"Updated description\"}'"
echo "# Test 4.3: Update both name and description"
make_request "POST" "/collections/${ALL_COLLECTION_IDS[0]}" \
"-d '{\"name\":\"Updated Name\",\"description\":\"Updated description\"}'"
echo "# Test 4.4: Generate description"
make_request "POST" "/collections/${ALL_COLLECTION_IDS[0]}" \
"-d '{\"generate_description\":true}'"
echo "# Test 4.5: Invalid update (both description and generate_description)"
make_request "POST" "/collections/${ALL_COLLECTION_IDS[0]}" \
"-d '{\"description\":\"New description\",\"generate_description\":true}'"
echo "# 5. Delete Collection Tests"
echo "=== Testing Collection Deletion ==="
echo "# Test 5.1: Delete existing collection"
make_request "DELETE" "/collections/${ALL_COLLECTION_IDS[0]}" ""
echo "# Test 5.2: Delete non-existent collection"
make_request "DELETE" "/collections/00000000-0000-0000-0000-000000000000" ""
echo "# 6. Document Management Tests"
echo "=== Testing Document Management ==="
echo "# Test 6.1: Add document to collection"
make_request "POST" "/collections/${ALL_COLLECTION_IDS[0]}/documents/456e789a-b12c-34d5-e678-901234567890" ""
echo "# Test 6.2: List documents in collection"
make_request "GET" "/collections/${ALL_COLLECTION_IDS[0]}/documents?offset=0&limit=10" ""
echo "# Test 6.3: List documents with invalid pagination"
make_request "GET" "/collections/${ALL_COLLECTION_IDS[0]}/documents?offset=-1&limit=10" ""
make_request "GET" "/collections/${ALL_COLLECTION_IDS[0]}/documents?offset=0&limit=1001" ""
echo "# Test 6.4: Remove document from collection"
make_request "DELETE" "/collections/${ALL_COLLECTION_IDS[0]}/documents/456e789a-b12c-34d5-e678-901234567890" ""
echo "# 7. User Management Tests"
echo "=== Testing User Management ==="
echo "# Test 7.1: List users in collection"
make_request "GET" "/collections/${ALL_COLLECTION_IDS[0]}/users?offset=0&limit=10" ""
echo "# Test 7.2: List users with invalid pagination"
make_request "GET" "/collections/${ALL_COLLECTION_IDS[0]}/users?offset=-1&limit=10" ""
make_request "GET" "/collections/${ALL_COLLECTION_IDS[0]}/users?offset=0&limit=1001" ""
echo "# Test 7.3: Add user to collection"
make_request "POST" "/collections/${ALL_COLLECTION_IDS[0]}/users/789a012b-c34d-5e6f-a789-012345678901" ""
echo "# Test 7.4: Remove user from collection"
make_request "DELETE" "/collections/${ALL_COLLECTION_IDS[0]}/users/789a012b-c34d-5e6f-a789-012345678901" ""
echo "# 8. Entity Extraction Tests"
echo "=== Testing Entity Extraction ==="
echo "# Test 8.1: Extract entities with default settings"
make_request "POST" "/collections/${ALL_COLLECTION_IDS[0]}/extract" \
"-d '{\"run_type\":\"RUN\"}'"
echo "# Test 8.2: Extract entities with custom settings"
make_request "POST" "/collections/${ALL_COLLECTION_IDS[0]}/extract" \
"-d '{\"run_type\":\"RUN\",\"settings\":{\"model\":\"gpt-4\"}}'"
echo "# Test 8.3: Extract entities without orchestration"
make_request "POST" "/collections/${ALL_COLLECTION_IDS[0]}/extract" \
"-d '{\"run_type\":\"RUN\",\"run_with_orchestration\":false}'"
echo "# Test 8.4: Get estimate only"
make_request "POST" "/collections/${ALL_COLLECTION_IDS[0]}/extract" \
"-d '{\"run_type\":\"ESTIMATE\"}'"
File diff suppressed because one or more lines are too long
@@ -1,42 +0,0 @@
On the first part of the journey
I was looking at all the life
There were plants and birds and rocks and things
There was sand and hills and rings
The first thing I met was a fly with a buzz
And the sky with no clouds
The heat was hot and the ground was dry
But the air was full of sound
I've been through the desert on a horse with no name
It felt good to be out of the rain
In the desert you can remember your name
'Cause there ain't no one for to give you no pain
La, la ...
After two days in the desert sun
My skin began to turn red
After three days in the desert fun
I was looking at a river bed
And the story it told of a river that flowed
Made me sad to think it was dead
You see I've been through the desert on a horse with no name
It felt good to be out of the rain
In the desert you can remember your name
'Cause there ain't no one for to give you no pain
La, la ...
After nine days I let the horse run free
'Cause the desert had turned to sea
There were plants and birds and rocks and things
there was sand and hills and rings
The ocean is a desert with it's life underground
And a perfect disguise above
Under the cities lies a heart made of ground
But the humans will give no love
You see I've been through the desert on a horse with no name
It felt good to be out of the rain
In the desert you can remember your name
'Cause there ain't no one for to give you no pain
La, la ...
Binary file not shown.

Before

Width:  |  Height:  |  Size: 201 KiB

Binary file not shown.
@@ -1 +0,0 @@
[{"@id": "data:text/plain;charset=US-ASCII,instituto%20superior%20tecnico", "@type": "http://www.cleverthis.com/university#Organization"}, {"@id": "data:text/plain;charset=US-ASCII,lisbon", "@type": "http://www.cleverthis.com/university#City"}, {"@id": "data:text/plain;charset=US-ASCII,lisbon", "@type": "http://www.cleverthis.com/university#Campus"}, {"@id": "data:text/plain;charset=US-ASCII,loures", "@type": "http://www.cleverthis.com/university#Campus"}, {"@id": "data:text/plain;charset=US-ASCII,Rogerio%20colaco", "@type": "http://www.cleverthis.com/university#Person"}, {"@id": "data:text/plain;charset=US-ASCII,1049-001%20lisboa", "@type": "data:,"}, {"@id": "data:text/plain;charset=US-ASCII,Instituto%20superior%20tecnico", "@type": "http://www.cleverthis.com/university#University", "http://www.cleverthis.com/university#campus": {"@id": "data:text/plain;charset=US-ASCII,three%20campuses", "@type": "http://www.cleverthis.com/university#Campus"}, "data:,": {"@id": "data:text/plain;charset=US-ASCII,tecnico", "@type": "http://www.cleverthis.com/university#Name"}, "http://www.cleverthis.com/university#country": {"@id": "data:text/plain;charset=US-ASCII,portugal", "@type": "http://www.cleverthis.com/university#Country"}}, {"@id": "data:text/plain;charset=US-ASCII,taguspark", "@type": "http://www.cleverthis.com/university#Campus"}, {"@id": "data:text/plain;charset=US-ASCII,oeiras", "@type": "http://www.cleverthis.com/university#City"}, {"@id": "data:text/plain;charset=US-ASCII,alameda", "@type": "http://www.cleverthis.com/university#Campus"}]
-180
View File
@@ -1,180 +0,0 @@
#!/bin/bash
. ./test_shared_functions.sh
JSON_FILE=./test_data/Instituto_Superior_Tecnico.md.json
HTML_FILE=./test_data/AMERICA-VENTURE_HIGHWAY.html
MP3_FILE=./test_data/flight_of_the_bumblebee_2.mp3
JPG_FILE=./test_data/China_Beijing_Summer_Palace2.jpg
EXCEL_FILE=./test_data/Worldbank-Datasheet-02082022.xlsx
TEXT_FILE=./test_data/America-Horse_with_no_name.txt
OPEN_OFFICE_FILE=./test_data/GAR_ARchitecture.odg
echo "# 1. Create Document Tests"
echo "=== Testing Document Creation ==="
echo "# Test 1.1: Create document with file, use different formats"
make_request "POST" "/documents" \
"-F file=@${TEXT_FILE}" \
"-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_FILE=${DOCUMENT_ID}
make_request "POST" "/documents" \
"-F file=@${JSON_FILE}" \
"-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_JSON_FILE=${DOCUMENT_ID}
make_request "POST" "/documents" \
"-F file=@${HTML_FILE}" \
"-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_HTML_FILE=${DOCUMENT_ID}
make_request "POST" "/documents" \
"-F file=@${MP3_FILE}" \
"-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_MP3_FILE=${DOCUMENT_ID}
make_request "POST" "/documents" \
"-F file=@${JPG_FILE}" \
"-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_JPG_FILE=${DOCUMENT_ID}
make_request "POST" "/documents" \
"-F file=@${EXCEL_FILE}" \
"-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_EXCEL_FILE=${DOCUMENT_ID}
make_request "POST" "/documents" \
"-F file=@${OPEN_OFFICE_FILE}" \
"-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_OPEN_OFFICE_FILE=${DOCUMENT_ID}
echo "# Test 1.2: Create document with raw text"
make_request "POST" "/documents" \
'-F raw_text="On the first part of the journey
I was looking at all the life
There were plants and birds and rocks and things
There was sand and hills and rings
The first thing I met was a fly with a buzz
And the sky with no clouds
The heat was hot and the ground was dry
But the air was full of sound"' \
"-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_RAW_TEXT=${DOCUMENT_ID}
echo "# Test 1.3: Create document with chunks"
make_request "POST" "/documents" \
"-F chunks=\"[\\\"I've been through the desert on a horse with no name\\\",\\\"It felt good to be out of the rain\\\"]\"" \
"-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_CHUNKS=${DOCUMENT_ID}
echo "# Test 1.4: Create document with metadata"
make_request "POST" "/documents" \
"-F file=@${TEXT_FILE} -F metadata='{\"title\":\"Horse with No Name\",\"author\":\"America\"}'" \
"-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_FILE_WITH_METADATA=${DOCUMENT_ID}
echo "# Test 1.5: Create document with custom ingestion config"
make_request "POST" "/documents" \
"-F file=@${TEXT_FILE} -F ingestion_mode=hi-res -F ingestion_config='{\"chunk_size\":1000}'" \
"-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_FILE_WITH_CUSTOM_INGEST=${DOCUMENT_ID}
echo "# 2. List Documents Tests"
echo "=== Testing Document Listing ==="
echo "# Test 2.1: List all documents"
make_request "GET" "/documents" ""
echo "# Test 2.2: List documents with pagination"
make_request "GET" "/documents?offset=0&limit=10" ""
echo "# Test 2.3: List specific documents by IDs"
make_request "GET" "/documents?ids=${DOCUMENT_ID_FROM_CHUNKS},${DOCUMENT_ID_FROM_FILE_WITH_METADATA}" ""
echo "# Test 2.4: List documents with summary embeddings"
make_request "GET" "/documents?include_summary_embeddings=1" ""
echo "# 3. Get Document Tests"
echo "=== Testing Document Retrieval ==="
echo "# Test 3.1: Get document by ID"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}" ""
echo "# Test 3.2: Get non-existent document"
make_request "GET" "/documents/00000000-0000-0000-0000-000000000000" ""
echo "# 4. List Document Chunks Tests"
echo "=== Testing Document Chunks ==="
echo "# Test 4.1: List chunks with pagination"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_CHUNKS}/chunks?offset=0&limit=10" ""
echo "# Test 4.2: List chunks with vectors"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_CHUNKS}/chunks?include_vectors=true" ""
echo "# 5. Download Document Tests"
echo "=== Testing Document Download ==="
echo "# Test 5.1: Download document"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_FILE}/download" ""
echo "# 6. Delete Document Tests"
echo "=== Testing Document Deletion ==="
echo "# Test 6.1: Delete document by ID"
make_request "DELETE" "/documents/${DOCUMENT_ID_FROM_RAW_TEXT}" ""
echo "# Test 6.2: Delete documents by filter"
make_request "DELETE" "/documents/by-filter" \
"-d '{\"document_type\":{\"\$eq\":\"txt\"}}'"
echo "# 7. Document Collections Tests"
echo "=== Testing Document Collections ==="
echo "# Test 7.1: List document collections"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_CHUNKS}/collections?offset=0&limit=10" ""
echo "# 8. Entity Extraction Tests"
echo "=== Testing Entity Extraction ==="
echo "# Test 8.1: Extract entities with default settings"
make_request "POST" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}/extract" \
"-d '{\"run_type\":\"RUN\"}'"
echo "# Test 8.2: Extract entities with custom settings"
make_request "POST" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}/extract" \
"-d '{\"run_type\":\"RUN\",\"settings\":{\"model\":\"gpt-4\"}}'"
echo "# 9. Entity Retrieval Tests"
echo "=== Testing Entity Retrieval ==="
echo "# Test 9.1: Get entities with pagination"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}/entities?offset=0&limit=10" ""
echo "# Test 9.2: Get entities with embeddings"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}/entities?include_embeddings=true" ""
echo "# 10. Relationship Tests"
echo "=== Testing Relationships ==="
echo "# Test 10.1: Get relationships with pagination"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}/relationships?offset=0&limit=10" ""
echo "# Test 10.2: Get relationships with filters"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}/relationships?entity_names=person,company&relationship_types=works_for" ""
echo "# 11. Document Search Tests"
echo "=== Testing Document Search ==="
echo "# Test 11.1: Basic search"
make_request "POST" "/documents/search" \
"-d '{\"query\":\"test query\",\"search_mode\":\"basic\"}'"
echo "# Test 11.2: Advanced search with custom settings"
make_request "POST" "/documents/search" \
"-d '{\"query\":\"test query\",\"search_mode\":\"advanced\",\"search_settings\":{\"filters\":{\"document_type\":\"pdf\"}}}'"
echo "# Test 11.3: Custom search with complex filters"
make_request "POST" "/documents/search" \
"-d '{\"query\":\"test query\",\"search_mode\":\"custom\",\"search_settings\":{\"filters\":{\"created_at\":{\"\$gt\":\"2023-01-01\"}}}}'"
-98
View File
@@ -1,98 +0,0 @@
#!/bin/bash
#
# Generates random valid UUID
#
generate_uuid() {
local N B C='89ab'
local RES=''
for (( N=0; N < 16; ++N )) do
B=$(( $RANDOM % 256 ))
case $N in
6)
RES="${RES}$(printf '4%x' $(( B % 16 )))"
;;
8)
RES="${RES}$(printf '%c%x' ${C:$RANDOM%${#C}:1} $(( B % 16 )))"
;;
3 | 5 | 7 | 9)
RES="${RES}$(printf '%02x-' $B)"
;;
*)
RES="${RES}$(printf '%02x' $B)"
;;
esac
done
echo -n $RES
}
#
# Extract document_id (if present) and HTTP response code from curl response passed as argument
#
extract_document_id() {
DOCUMENT_ID=$(echo $1 | grep document_id | tail -1 | sed -e 's/^.*"document_id": *"\([^"]*\)".*$/\1/')
HTTP_RETURN_CODE=$(echo echo $1 | tail -1 | sed -e 's/^.*\([0-9][0-9][0-9]\)$/\1/')
}
# Helper function to make curl requests
make_request() {
local method=$1
local endpoint=$2
local data=$3
local headers=$4
if [ "xx${headers}" == "xx" ]; then
headers='-H "Content-Type: application/json"'
fi
# echo "Testing $method $endpoint"
# echo "Request data: $data"
# echo "Request headers: $headers"
CMD="curl -X $method \"$API_BASE$endpoint\" -s -w \"\n%{http_code}\" -H \"Authorization: Bearer $AUTH_TOKEN\" $headers $data"
echo $CMD
RES=$(eval $CMD 2>/dev/null)
echo "Response: [${RES}]"
extract_document_id "$RES"
# echo "document_id --> ${DOCUMENT_ID}"
if [[ ${HTTP_RETURN_CODE} -lt 300 ]]; then
echo "SUCCESS. code = ${HTTP_RETURN_CODE}"
else
echo "ERROR. code = ${HTTP_RETURN_CODE}"
fi
echo -e "\n\n"
if [ "${EXECUTE_ALL}" != "-a" ]; then
echo "Press any key to continue..."
read -n 1 -s -r
echo -e "\n\n"
fi
}
#
# Simple scanner of cmd line arguments, inferring meaning of the value
# It recognized:
# a) host URL (starts with 'http' prefix)
# b) JWT token (starts with 'ey' prefix)
# c) flag -a to execute all tests w/o waiting for user pressing any key
#
while [ $# -gt 0 ]; do
echo "Argument: [${1}]"
if [[ "$1" == ey* ]]; then
AUTH_TOKEN=$1
fi
if [[ "$1" == http* ]]; then
HOST=$1
fi
if [[ "$1" == "-a" ]]; then
EXECUTE_ALL=$1
fi
shift # Removes $1 and shifts remaining args left
done
if [ "xx${HOST}" == "xx" ]; then
HOST="http://localhost:7272"
fi
API_BASE="$HOST/v3"
# Once integrated with CleverMicro, token can be obtained from CleverMicro
# if [ "xx${AUTH_TOKEN}" == "xx" ]; then
# AUTH_TOKEN=$(curl $HOST/api/v0/login -s -d '{"username":"hugo@cleverthis.com","password":"tellmeall"}' -H "Content-Type: application/json" | sed -e "s/.*access_token...\(.*\)..$/\1/")
# fi
+256 -35
View File
@@ -5,51 +5,129 @@ It is NOT MEANT to be used as a standalone module. It WILL NOT run outside the C
"""
import importlib
import pathlib
import re
from collections import defaultdict
from threading import Thread
from typing import Any, List, Optional
from typing import Optional, List, Dict
from aiohttp import ClientSession
from fastapi.routing import APIRoute
from fastapi.security import OAuth2PasswordRequestForm
from starlette.datastructures import UploadFile, Headers
# CleverSwarm specific imports
import core.deps
from fastapi.routing import APIRoute
from rest.v0 import api
from schemas.user_schema import UserSchema
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter
from amqp.adapter.data_parser import AMQDataParser
# AMQP specific imports
from amqp.adapter.data_response_factory import AMQResponseFactory
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter, AMQMessage
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQResponse
from amqp.service.amq_service import AMQService
# Suffix for the method name with serialization wrapper, prefer use this one over the orig endpoint
PREFERRED_SUFFIX = "_json"
# =================================================================================================
# =========================== S u p p o r t i n g f u n c t i o n s =============================
# =================================================================================================
def get_callable_for_name(module_name: str, function_name: str) -> Optional[callable]:
def clone_and_adapt_route(original_route: APIRoute) -> APIRoute:
"""
Get the function/Callable for the module_name.function_name if it exists in the module, else return None.
Creates a clone of an APIRoute. It checks if there's preferred endpoint with json serialization.
It also adjusts the path regex to ensure it matches the path in different API versions.
"""
# Check if the route has a preferred endpoint with json serialization
_preferred_endpoint = get_prefered_endpoint(
getattr(original_route.endpoint, "__module__"),
original_route.endpoint.__name__ + PREFERRED_SUFFIX,
)
# Need to adapt also the path regex to ensure it matches the path in different API versions
# Ensures the leading '^' is followed by any match and there's no '/' before final '$'
_new_regex_string = (
"^.*" + original_route.path_regex.pattern.lstrip("^").rstrip("/$") + "$"
)
_new_route: APIRoute = APIRoute(
path=original_route.path,
endpoint=(
_preferred_endpoint if _preferred_endpoint else original_route.endpoint
),
response_model=original_route.response_model,
status_code=original_route.status_code,
tags=original_route.tags.copy() if original_route.tags else None,
dependencies=(
original_route.dependencies.copy() if original_route.dependencies else None
),
summary=original_route.summary,
description=original_route.description,
response_description=original_route.response_description,
responses=original_route.responses.copy() if original_route.responses else None,
deprecated=original_route.deprecated,
name=original_route.name,
methods=original_route.methods.copy() if original_route.methods else None,
operation_id=original_route.operation_id,
response_model_include=original_route.response_model_include,
response_model_exclude=original_route.response_model_exclude,
response_model_by_alias=original_route.response_model_by_alias,
response_model_exclude_unset=original_route.response_model_exclude_unset,
response_model_exclude_defaults=original_route.response_model_exclude_defaults,
response_model_exclude_none=original_route.response_model_exclude_none,
include_in_schema=original_route.include_in_schema,
response_class=original_route.response_class,
dependency_overrides_provider=original_route.dependency_overrides_provider,
callbacks=original_route.callbacks.copy() if original_route.callbacks else None,
openapi_extra=(
original_route.openapi_extra.copy()
if original_route.openapi_extra
else None
),
generate_unique_id_function=original_route.generate_unique_id_function,
)
_new_route.path_regex = re.compile(_new_regex_string)
return _new_route
def get_prefered_endpoint(module_name: str, function_name: str) -> Optional[callable]:
"""Get the function if it exists in the module, else return None."""
try:
_module = importlib.import_module(module_name)
_preferred_endpoint = getattr(_module, function_name)
return _preferred_endpoint if callable(_preferred_endpoint) else None
_prefered_endpoint = getattr(_module, function_name)
return _prefered_endpoint if callable(_prefered_endpoint) else None
except Exception:
return None
# =================================================================================================
# ======================== C l e v e r S w a r m A M Q A d a p t e r ========================
# =================================================================================================
def group_routes_by_method(api_routes: List[APIRoute]) -> Dict[str, List[APIRoute]]:
"""
Groups API routes by their HTTP method. This is important to ensure that the URLPath is matched
correctly, as different methods may share the same URLPath (e.g. GET, POST PUT, DELETE).
"""
_method_to_routes = defaultdict(
list
) # Automatically initializes new lists for new keys
for _route in api_routes:
if (
_route.methods
): # Ensure methods exist (should always be true for valid routes)
for method in _route.methods:
_method_to_routes[method].append(clone_and_adapt_route(_route))
return dict(_method_to_routes)
class CleverSwarmAmqpAdapter(CleverThisServiceAdapter):
"""
CleverSwarm AMQP Adapter for CleverSwarm API. Implements the CleverSwarm specific logic that overrides
common logic defined by AMQ Adapter library. The specific logic here is:
1. Extract subject from JWT token and convert it to UserSchema to avail authorized user argument
2. Find JSON wrapper for endpoint Callable, and if exists, use it instead of the original endpoint.
CleverSwarm AMQP Adapter for CleverSwarm API. Implements the interface to CleverSwarm API.
"""
def __init__(self, amq_configuration: AMQConfiguration, routes: List[APIRoute]):
super().__init__(amq_configuration, routes=routes, session=None)
def __init__(
self, amq_configuration: AMQConfiguration, session: ClientSession = None
):
super().__init__(amq_configuration, session)
self.routes = group_routes_by_method(api_routes=api.api_router.routes)
_amq_service: AMQService = AMQService(amq_configuration, self)
self.thread = Thread(target=_amq_service.run)
self.thread.start()
@@ -57,25 +135,168 @@ class CleverSwarmAmqpAdapter(CleverThisServiceAdapter):
async def get_current_user(self, token: str) -> UserSchema:
"""
Overrides the default get_current_user method to use the token from the AMQP message.
In CleverSwarm context, it means simply to call provided get_current_user function.
:param token: JWT token from the AMQP message
:return: UserSchema object
"""
return await core.deps.get_current_user(token)
def get_endpoint(self, endpoint: Any) -> Any:
"""
If the specific endpoint has a matching JSON wrapper (identified as {endpoint.name}_json ),
then use it as preferred way to invoke the endpoint.
# ==================================================================================================
# ===================== C l e v e r S w a r m A P I m e t h o d s ============================
# ==================================================================================================
:param endpoint: the Callable for the endpoint
:return: new Callable representing the JSON wrapper endpoint, if it exists or the original Callable.
async def get(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
"""
_preferred_endpoint = get_callable_for_name(
getattr(endpoint, "__module__"),
endpoint.__name__ + PREFERRED_SUFFIX,
Handles the GET requests for the CleverSwarm API.
:param message: message from the AMQP that contains the GET request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: AMQResponse to be sent back to the API Gateway via AMQP
"""
return await self.handle_no_body_payload(message, auth_user, "GET")
async def put(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
"""
Handles the PUT requests for the CleverSwarm API.
:param message: message from the AMQP that contains the PUT request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: AMQResponse to be sent back to the API Gateway via AMQP
"""
return await self.handle_body_payload(message, auth_user, "PUT")
async def post(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
"""
Handles the POST requests for the CleverSwarm API.
:param message: message from the AMQP that contains the POST request, including body and headers
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: AMQResponse to be sent back to the API Gateway via AMQP
"""
return await self.handle_body_payload(message, auth_user, "POST")
async def head(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
"""
Handles the HEAD requests for the CleverSwarm API.
"""
return await self.handle_no_body_payload(message, auth_user, "HEAD")
async def delete(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
"""
Handles the DELETE requests for the CleverSwarm API.
"""
return await self.handle_no_body_payload(message, auth_user, "DELETE")
async def options(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
"""
Handles the OPTIONS requests for the CleverSwarm API.
"""
return await self.handle_no_body_payload(message, auth_user, "OPTIONS")
async def patch(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
"""
Handles the PATCH requests for the CleverSwarm API.
"""
return await self.handle_no_body_payload(message, auth_user, "PATCH")
async def handle_no_body_payload(
self, message: AMQMessage, auth_user: UserSchema, method: str
) -> AMQResponse:
"""
Handles the requests that do not have a body payload.
:param message: message from the AMQP that contains the request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: AMQResponse to be sent back to the API Gateway via AMQP
"""
_routes = self.routes.get(method, [])
_path, _args = AMQDataParser.parse_get_url(message.path)
for _route in _routes:
_match = _route.path_regex.match(_path)
if _match:
_args.update(_match.groupdict())
_args["authenticated_user"] = auth_user
# If the path matches, call the corresponding function
return await CleverThisServiceAdapter.call_cleverthis_api(
message,
_args,
api_method=_route.endpoint,
success_code=_route.status_code,
loop=self.loop,
)
return AMQResponseFactory.of(
message.id,
404,
"text/plain",
str(f"Destination location [{message.path}] does not exists").encode(
"utf-8"
),
message.trace_info,
)
async def handle_body_payload(
self, message: AMQMessage, auth_user: UserSchema, method: str
) -> AMQResponse:
"""
Handles the requests that have a body payload, including Multipart data.
:param message: message from the AMQP that contains the request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:param method: HTTP method (POST, PUT, etc.)
:return: AMQResponse to be sent back to the API Gateway via AMQP
"""
_form_data = AMQDataParser.parse_request_body(message)
_routes = self.routes.get(method, [])
_path, _query_args = AMQDataParser.parse_get_url(message.path)
for _route in _routes:
_match = _route.path_regex.match(_path)
if _match:
# If the path matches, call the corresponding function
_form_data.update(_match.groupdict())
_form_data.update(_query_args)
for _body_param in _route.dependant.body_params:
if _body_param.name in _form_data:
_file_data = _form_data[_body_param.name]
actual_filepath = message.extra_data[_body_param.name]
_form_data[_body_param.name] = (
_body_param.type_.__name__ == "UploadFile"
and UploadFile(
file=pathlib.Path(actual_filepath).open("rb"),
size=_file_data["size"],
filename=_file_data["filename"],
headers=Headers(
{"Content-Type": _file_data["mediaType"]}
),
)
or _form_data[_body_param.name]
)
_verified_args = {}
_required_args = getattr(_route.endpoint, "__annotations__", {})
for arg in _required_args:
if arg in _form_data:
_verified_args[arg] = _form_data[arg]
else:
# needed for /login input
if arg == "authenticated_user":
_verified_args["authenticated_user"] = auth_user
if _required_args[arg].__name__ == "OAuth2PasswordRequestForm":
try:
_verified_args[arg] = OAuth2PasswordRequestForm(
**_form_data
)
except Exception as e:
print(f"Error: {e}")
return await CleverThisServiceAdapter.call_cleverthis_api(
message,
_verified_args,
api_method=_route.endpoint,
success_code=_route.status_code if _route.status_code else 200,
loop=self.loop,
)
return AMQResponseFactory.of(
message.id,
404,
"text/plain",
str(f"Destination location [{message.path}] does not exists").encode(
"utf-8"
),
message.trace_info,
)
return _preferred_endpoint if _preferred_endpoint is not None else endpoint
+2 -2
View File
@@ -1,7 +1,7 @@
from clever_swarm_adapter import CleverSwarmAmqpAdapter
from amqp.config.amq_configuration import AMQConfiguration
from clever_swarm_adapter import CleverSwarmAmqpAdapter
if __name__ == "__main__":
adapter = CleverSwarmAmqpAdapter(AMQConfiguration("./application.properties.local"))
print("Press ^C to exit...")
+8 -20
View File
@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
[project]
name = "amq_adapter"
version = "0.2.33"
description = "The Python implementation of the AMQ Adapter for CleverMicro"
version = "0.2.23-dev"
description = "The Python implementation of the AMQ Adaptor for CleverMicro"
readme = "README.md"
authors = [
{ name = "Stanislav Hejny", email = "stanislav.hejny@cleverthis.com" }
@@ -19,16 +19,15 @@ classifiers = [
]
requires-python = ">=3.11"
dependencies = [
"aio-pika",
"aiohttp",
"fastapi",
"aio-pika~=9.5",
"aiohttp~=3.11",
"aio_pika~=9.5",
"fastapi~=0.114",
"opentelemetry-api",
"opentelemetry-sdk",
"opentelemetry-exporter-otlp",
"opentelemetry-exporter-prometheus",
"opentelemetry-instrumentation-pika",
"pika",
"psutil",
"python_multipart"
]
@@ -36,22 +35,11 @@ dependencies = [
where = ["."]
include = ["amqp*"]
[tool.black]
line-length = 100
[tool.pytest.ini_options]
asyncio_default_fixture_loop_scope = "module"
[project.optional-dependencies]
dev = [
"pytest>=7.0",
"pytest-cov>=4.0",
"pytest-asyncio",
"build",
"twine",
"black",
"isort",
"flake8",
"pre-commit",
"pylint"
"twine"
]
View File
View File
-512
View File
@@ -1,512 +0,0 @@
#
# 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 == mock_config.backpressure.threshold_threads
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()
+72 -276
View File
@@ -1,105 +1,14 @@
import asyncio
import unittest
from enum import Enum
from typing import Any, List, Optional
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 pydantic import BaseModel
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.adapter.cleverthis_service_adapter import CleverThisServiceAdapter, AMQMessage
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import DataResponse
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE
from amqp.model.model import DataMessage, DataResponse
class NestedModel(BaseModel):
id: int
name: str
class ComplexModel(BaseModel):
nested: NestedModel
items: List[str]
class TestEnum(Enum):
VALUE1 = "value1"
VALUE2 = "value2"
class TestModel(BaseModel):
name: str
value: int
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
class TestCleverThisServiceAdapter(unittest.TestCase):
def setUp(self):
# Mock AMQConfiguration
self.mock_amq_config = AMQConfiguration("./application.properties.local")
@@ -109,13 +18,11 @@ class TestCleverThisServiceAdapter(unittest.IsolatedAsyncioTestCase):
# Create adapter instance
self.adapter = CleverThisServiceAdapter(
amq_configuration=self.mock_amq_config, routes=[], session=self.mock_session
amq_configuration=self.mock_amq_config, session=self.mock_session
)
self.adapter.service_host = "testhost"
self.adapter.authenticated_user_arg_name = "user"
def test_init(self):
self.assertEqual(self.adapter.service_host, "testhost")
self.assertEqual(self.adapter.service_host, "localhost")
self.assertEqual(self.adapter.service_port, 8080)
self.assertFalse(self.adapter.require_authenticated_user)
@@ -141,27 +48,26 @@ class TestCleverThisServiceAdapter(unittest.IsolatedAsyncioTestCase):
@patch.object(CleverThisServiceAdapter, "get_current_user")
async def test_on_message_unauthenticated(self, mock_get_user):
mock_message = MagicMock(spec=AMQMessage)
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 = {}
mock_message.headers = {}
mock_message.path = "/some/path"
mock_message.method = "GET"
mock_message.id = "123"
mock_message.trace_info = {}
# 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")
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: MagicMock, mock_get_user):
async def test_on_message_authenticated(self, mock_get, mock_get_user):
mock_message = MagicMock(spec=AMQMessage)
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_message.headers = {"Authorization": ["Bearer valid_token"]}
mock_message.path = "/some/path"
mock_message.method = "GET"
mock_message.id = "123"
mock_message.trace_info = {}
mock_user = MagicMock()
mock_get_user.return_value = mock_user
@@ -169,7 +75,6 @@ class TestCleverThisServiceAdapter(unittest.IsolatedAsyncioTestCase):
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)
@@ -179,14 +84,16 @@ class TestCleverThisServiceAdapter(unittest.IsolatedAsyncioTestCase):
@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.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_message.headers = {}
mock_message.path = "/login"
mock_message.method = "POST"
mock_message.id = "123"
mock_message.trace_info = {}
mock_response = MagicMock(spec=DataResponse)
with patch.object(self.adapter, "post", return_value=mock_response) as mock_post:
with patch.object(
self.adapter, "post", return_value=mock_response
) as mock_post:
response = await self.adapter.on_message(mock_message)
self.assertEqual(response, mock_response)
mock_get_user.assert_not_called()
@@ -203,11 +110,11 @@ class TestCleverThisServiceAdapter(unittest.IsolatedAsyncioTestCase):
for method in methods:
with self.subTest(method=method):
mock_message = MagicMock(spec=AMQMessage)
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 = {}
mock_message.headers = {"Authorization": ["Bearer valid_token"]}
mock_message.path = "/some/path"
mock_message.method = method
mock_message.id = "123"
mock_message.trace_info = {}
with patch.object(
self.adapter, method.lower(), return_value=mock_response
@@ -217,11 +124,11 @@ class TestCleverThisServiceAdapter(unittest.IsolatedAsyncioTestCase):
mock_method.assert_called_once()
async def test_get_with_session(self):
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_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_response = MagicMock(spec=ClientResponse)
mock_response.status = 200
@@ -232,7 +139,7 @@ class TestCleverThisServiceAdapter(unittest.IsolatedAsyncioTestCase):
response = await self.adapter.get(mock_message, None)
self.assertEqual(response.response_code(), "200")
self.assertEqual(response.response_code, 200)
self.mock_session.get.assert_called_once_with(
"http://testhost:8080/api/test",
headers={"X-Test": "value", "trace-id": "123"},
@@ -241,50 +148,50 @@ class TestCleverThisServiceAdapter(unittest.IsolatedAsyncioTestCase):
async def test_get_without_session(self):
self.adapter.session = None
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 = {}
mock_message = MagicMock(spec=DataMessage)
mock_message.path = "/api/test"
mock_message.id = "msg123"
mock_message.trace_info = {}
response = await self.adapter.get(mock_message, None)
self.assertEqual(response.response_code(), "404")
self.assertEqual(response.response_code, 404)
async def test_put_with_form_handling(self):
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_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_response: DataResponse = MagicMock(spec=DataResponse)
mock_response.response_code.return_value = "200"
mock_response.content_type.return_value = "application/json"
mock_response = MagicMock(spec=ClientResponse)
mock_response.status = 200
mock_response.content_type = "application/json"
mock_response.read = AsyncMock(return_value=b'{"result": "success"}')
self.mock_session.put.return_value = mock_response
with patch.object(
self.adapter,
"_authorize_forward_request",
wraps=self.adapter._authorize_forward_request,
"handle_possible_form",
wraps=self.adapter.handle_possible_form,
) as mock_handler:
response = await self.adapter.put(mock_message, None)
self.assertEqual(response.response_code(), "200")
self.assertEqual(response.response_code, 200)
mock_handler.assert_called_once()
async def test_post_with_form_handling(self):
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_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_response: DataResponse = MagicMock(spec=DataResponse)
mock_response.response_code.return_value = "201"
mock_response = MagicMock(spec=ClientResponse)
mock_response.status = 201
mock_response.content_type = "application/json"
mock_response.read = AsyncMock(return_value=b'{"result": "created"}')
@@ -292,139 +199,28 @@ class TestCleverThisServiceAdapter(unittest.IsolatedAsyncioTestCase):
with patch.object(
self.adapter,
"_authorize_forward_request",
wraps=self.adapter._authorize_forward_request,
"handle_possible_form",
wraps=self.adapter.handle_possible_form,
) as mock_handler:
response = await self.adapter.post(mock_message, None)
self.assertEqual(response.response_code(), "201")
self.assertEqual(response.response_code, 201)
mock_handler.assert_called_once()
async def test_head_not_implemented(self):
mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message = MagicMock(spec=DataMessage)
mock_message.id = "msg123"
with self.assertRaises(NotImplementedError):
await self.adapter.head(mock_message, None)
def test_basic_type_conversion(self):
async def test_method(x: int, y: float, z: bool, text: str):
pass
async def test_delete_not_implemented(self):
mock_message = MagicMock(spec=DataMessage)
mock_message.id = "msg123"
args = {"x": "123", "y": "45.67", "z": "true", "text": "hello"}
result = self.adapter._verify_and_instantiate_arguments(args, test_method)
self.assertEqual(result["x"], 123)
self.assertEqual(result["y"], 45.67)
self.assertEqual(result["z"], True)
self.assertEqual(result["text"], "hello")
def test_enum_conversion(self):
async def test_method(status: TestEnum):
pass
args = {"status": "value1"}
result = self.adapter._verify_and_instantiate_arguments(args, test_method)
self.assertEqual(result["status"], TestEnum.VALUE1)
def test_optional_parameter(self):
async def test_method(required: str, optional: Optional[str] = None):
pass
args = {"required": "test"}
result = self.adapter._verify_and_instantiate_arguments(args, test_method)
self.assertEqual(result["required"], "test")
self.assertIsNone(result["optional"])
def test_list_parameter(self):
async def test_method(items: List[str]):
pass
args = {"items": "item1,item2,item3"}
result = self.adapter._verify_and_instantiate_arguments(args, test_method)
self.assertEqual(result["items"], ["item1", "item2", "item3"])
def test_pydantic_model_parameter(self):
async def test_method(model: TestModel):
pass
args = {"model": '{"name": "test", "value": 42}'}
result = self.adapter._verify_and_instantiate_arguments(args, test_method)
self.assertIsInstance(result["model"], TestModel)
self.assertEqual(result["model"].name, "test")
self.assertEqual(result["model"].value, 42)
def test_authenticated_user_parameter(self):
async def test_method(user: Any):
pass
user_obj = MagicMock()
args = {"user": user_obj}
result = self.adapter._verify_and_instantiate_arguments(args, test_method)
self.assertEqual(result["user"], user_obj)
def test_missing_required_parameter(self):
async def test_method(required: str):
pass
args = {}
result = self.adapter._verify_and_instantiate_arguments(args, test_method)
self.assertNotIn("required", result)
def test_invalid_enum_value(self):
async def test_method(status: TestEnum):
pass
args = {"status": "invalid_value"}
result = self.adapter._verify_and_instantiate_arguments(args, test_method)
self.assertEqual(result["status"], "invalid_value")
def test_invalid_type_conversion(self):
async def test_method(x: int):
pass
args = {"x": "not_a_number"}
result = {}
try:
result = self.adapter._verify_and_instantiate_arguments(args, test_method)
except Exception as e:
self.assertTrue(isinstance(e, TypeError))
self.assertEqual(len(result), 0)
def test_complex_nested_model(self):
async def test_method(model: ComplexModel):
pass
args = {"model": {"nested": {"id": 1, "name": "test"}, "items": ["item1", "item2"]}}
result = self.adapter._verify_and_instantiate_arguments(args, test_method)
self.assertIsInstance(result["model"], ComplexModel)
self.assertEqual(result["model"].nested.id, 1)
self.assertEqual(result["model"].nested.name, "test")
self.assertEqual(result["model"].items, ["item1", "item2"])
def test_default_factory_parameter(self):
class DefaultFactory:
@staticmethod
def default_factory():
return "default_value"
async def test_method(param: str = DefaultFactory()):
pass
args = {}
result = self.adapter._verify_and_instantiate_arguments(args, test_method)
self.assertEqual(result["param"], "default_value")
# Assuming delete is not implemented in base class
with self.assertRaises(NotImplementedError):
await self.adapter.delete(mock_message, None)
if __name__ == "__main__":
+18 -12
View File
@@ -10,21 +10,25 @@ class DataMessageFactoryTest(TestCase):
msg = DataMessageFactory.from_bytes(raw.encode("utf-8"))
self.assertIsNotNone(msg.id)
self.assertEqual(
msg.id().as_string(),
msg.id.as_string(),
"64F3B5AF4D00000A4000".lower(),
"SnowflakeID incorrectly parsed.",
)
self.assertEqual(
msg.headers()["Content-Type"][0],
msg.headers["Content-Type"][0],
"multipart/form-data; boundary=------------------------j1RjIqLrFH07XMsyMvpS1v",
)
def test_generate_snowflake_id(self):
_id = DataMessageFactory.get_instance(1).generate_snowflake_id()
self.assertEqual("00001000", _id.as_string()[-8:], "Instance ID not set correctly")
self.assertEqual(
"00001000", _id.as_string()[-8:], "Instance ID not set correctly"
)
_id = DataMessageFactory(2).generate_snowflake_id()
self.assertEqual("00002000", _id.as_string()[-8:], "Instance ID not set correctly")
self.assertEqual(
"00002000", _id.as_string()[-8:], "Instance ID not set correctly"
)
_f = DataMessageFactory(127)
_i0 = _f.generate_snowflake_id()
@@ -61,9 +65,9 @@ class DataMessageFactoryTest(TestCase):
{"k1": ["v1"], "k2": ["v2"]},
"Quick brown fox",
)
self.assertEqual("localhost", _req.domain(), "Domain not set correctly")
self.assertEqual("/test/me", _req.path(), "Context path not set correctly")
self.assertEqual("v1", _req.headers()["k1"][0], "Headers not set correctly")
self.assertEqual("localhost", _req.domain, "Domain not set correctly")
self.assertEqual("/test/me", _req.path, "Context path not set correctly")
self.assertEqual("v1", _req.headers["k1"][0], "Headers not set correctly")
self.assertEqual(b"Quick brown fox", _req.body(), "Body not correctly decoded")
def test_serialize(self):
@@ -78,11 +82,11 @@ class DataMessageFactoryTest(TestCase):
_ser = _f.serialize(_req)
self.assertEqual(65, _ser[2])
_deser = _f.from_bytes(_ser)
self.assertEqual(_req.id().as_string(), _deser.id().as_string())
self.assertEqual("localhost", _deser.domain())
self.assertEqual("/test/me", _deser.path())
self.assertEqual(_req.id.as_string(), _deser.id.as_string())
self.assertEqual("localhost", _deser.domain)
self.assertEqual("/test/me", _deser.path)
self.assertEqual(b"Quick brown fox", _deser.body())
self.assertEqual("v2", _deser.headers()["k2"][0])
self.assertEqual("v2", _deser.headers["k2"][0])
def test_amq_message_routing_key(self):
_f = DataMessageFactory(35)
@@ -120,4 +124,6 @@ class DataMessageFactoryTest(TestCase):
"Quick brown fox",
)
_as_str = _f.to_string(_req)
self.assertGreater(_as_str.index("domain='localhost'"), 0, "Expected content not found")
self.assertGreater(
_as_str.index("domain='localhost'"), 0, "Expected content not found"
)
-165
View File
@@ -1,165 +0,0 @@
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")
+6 -7
View File
@@ -2,7 +2,6 @@ from unittest import TestCase
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.data_response_factory import DataResponseFactory
from amqp.model.model import DataMessageMagicByte
from amqp.model.snowflake_id import SnowflakeId
@@ -12,8 +11,8 @@ class TestDataResponseFactory(TestCase):
_f = DataResponseFactory()
_g = DataMessageFactory(111)
_resp = _f.create_async_response_message(_g.generate_snowflake_id())
self.assertEqual("200", _resp.response_code())
self.assertEqual(b"OK", _resp.body())
self.assertEqual(200, _resp.response_code)
self.assertEqual(b"OK", _resp.response)
def test_serialize(self):
_f = DataResponseFactory()
@@ -23,7 +22,7 @@ class TestDataResponseFactory(TestCase):
_ser = _f.serialize(_resp)
self.assertEqual(
2,
str(_ser[2:]).index(DataMessageMagicByte.HTTP_REQUEST.value + _snowflakeId.as_string()),
str(_ser[2:]).index(_f.MAGIC_BYTE_HTTP_REQUEST + _snowflakeId.as_string()),
"Magic byte and SnowflakeId not found",
)
self.assertGreater(str(_ser[2:]).index("|e=|"), 2)
@@ -36,6 +35,6 @@ class TestDataResponseFactory(TestCase):
_resp = _f.create_async_response_message(_snowflakeId)
_ser = _f.serialize(_resp)
_deser = _f.from_bytes(_ser)
self.assertEqual(_snowflakeId.as_string(), _deser.id())
self.assertEqual("200", _deser.response_code())
self.assertEqual(b"OK", _deser.body())
self.assertEqual(_snowflakeId.as_string(), _deser.id)
self.assertEqual(200, _deser.response_code)
self.assertEqual(b"OK", _deser.response)
+2 -263
View File
@@ -1,22 +1,11 @@
import asyncio
import os
import tempfile
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 unittest import TestCase
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(IsolatedAsyncioTestCase):
class TestFileHandler(TestCase):
def test_ensure_directory_exists(self):
"""
Test function for ensure_directory_exists. Creates a temporary directory
@@ -59,253 +48,3 @@ class TestFileHandler(IsolatedAsyncioTestCase):
), 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)))
-138
View File
@@ -1,138 +0,0 @@
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
+17 -205
View File
@@ -1,226 +1,38 @@
import json
from datetime import date, datetime
from unittest import TestCase
from uuid import UUID
from pydantic import BaseModel
from amqp.adapter.pydantic_serializer import (
deserialize_object,
parse_url_query,
python_type_to_json,
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 = []
from amqp.adapter.pydantic_serializer import parse_url_query
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):
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)
assert True
def test_deserialize_object(self):
serialized = serialize_object(self.test_model)
deserialized = deserialize_object(serialized)
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__convert_value(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_uuid_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_datetime_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__is_int_string(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_serialize_to_json(self):
assert True
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"
)
self.assertEqual(path, "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741")
self.assertEqual(args, {"response_type": "JsonFile"})
assert path == "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741"
assert args == {"response_type": "JsonFile"}
path, args = parse_url_query(
"http://localhost:8080/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741"
)
self.assertEqual(path, "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741")
self.assertEqual(args, {})
class TestPythonTypeToJson:
"""
Test cases for the python_type_to_json function.
"""
def test_basic_types(self):
"""Test serialization of basic Python types."""
# Test dictionary
data = {"key": "value", "number": 42}
result = python_type_to_json(data)
assert json.loads(result) == data
# Test list
data = [1, 2, 3, "four"]
result = python_type_to_json(data)
assert json.loads(result) == data
# Test tuple
data = (1, 2, 3)
result = python_type_to_json(data)
assert json.loads(result) == [1, 2, 3] # Tuples are converted to lists
# Test set
data = {1, 2, 3}
result = python_type_to_json(data)
assert set(json.loads(result)) == data # Sets are converted to lists
def test_datetime_handling(self):
"""Test serialization of datetime objects."""
# Test datetime
test_datetime = datetime(2024, 1, 1, 12, 0, 0)
data = {"timestamp": test_datetime}
result = python_type_to_json(data)
assert json.loads(result) == {"timestamp": test_datetime.isoformat()}
# Test date
test_date = date(2024, 1, 1)
data = {"date": test_date}
result = python_type_to_json(data)
assert json.loads(result) == {"date": test_date.isoformat()}
def test_nested_structures(self):
"""Test serialization of nested data structures."""
data = {
"list_of_datetimes": [datetime(2024, 1, 1), datetime(2024, 1, 2)],
"dict_with_date": {"date": date(2024, 1, 1)},
"mixed": {"set": {1, 2, 3}, "tuple": (1, 2, 3), "datetime": datetime(2024, 1, 1)},
}
result = python_type_to_json(data)
parsed = json.loads(result)
assert parsed["list_of_datetimes"] == [d.isoformat() for d in data["list_of_datetimes"]]
assert parsed["dict_with_date"]["date"] == data["dict_with_date"]["date"].isoformat()
assert set(parsed["mixed"]["set"]) == data["mixed"]["set"]
assert parsed["mixed"]["tuple"] == list(data["mixed"]["tuple"])
assert parsed["mixed"]["datetime"] == data["mixed"]["datetime"].isoformat()
def test_custom_object_serialization(self):
"""Test serialization of custom objects with __dict__ attribute."""
class TestObject:
def __init__(self, value):
self.value = value
self.timestamp = datetime(2024, 1, 1)
obj = TestObject("test")
data = {"obj": obj}
result = python_type_to_json(data)
parsed = json.loads(result)
assert parsed["obj"]["value"] == "test"
assert parsed["obj"]["timestamp"] == obj.timestamp.isoformat()
def test_empty_structures(self):
"""Test serialization of empty data structures."""
# Empty dict
assert python_type_to_json({}) == "{}"
# Empty list
assert python_type_to_json([]) == "[]"
# Empty set
assert python_type_to_json(set()) == "[]"
# Empty tuple
assert python_type_to_json(()) == "[]"
def test_unicode_handling(self):
"""Test handling of Unicode characters."""
data = {"text": "Hello, 世界!"}
result = python_type_to_json(data)
assert json.loads(result) == data # ensure_ascii=False should preserve Unicode
assert path == "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741"
assert args == {}
-114
View File
@@ -1,114 +0,0 @@
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))
@@ -1,9 +1,9 @@
from unittest import TestCase
from amqp.adapter.service_message_factory import (
INVALID_MESSAGE,
RS,
ServiceMessageFactory,
RS,
INVALID_MESSAGE,
)
from amqp.model.model import ServiceMessage
from amqp.model.service_message_type import ServiceMessageType
@@ -13,11 +13,15 @@ class ServiceMessageTest(TestCase):
def test_to_string(self):
_f: ServiceMessageFactory = ServiceMessageFactory("bumble-bee")
message: ServiceMessage = _f.of(ServiceMessageType.ROUTING_DATA, "Duck", "Donald")
message: ServiceMessage = _f.of(
ServiceMessageType.ROUTING_DATA, "Duck", "Donald"
)
message_str = _f.to_string(message)
self.assertTrue("|type=01|" in message_str)
self.assertTrue("|body=Duck" in message_str)
message: ServiceMessage = _f.of(ServiceMessageType.ROUTING_DATA_REPLACE, "Duck", "Donald")
message: ServiceMessage = _f.of(
ServiceMessageType.ROUTING_DATA_REPLACE, "Duck", "Donald"
)
message_str = _f.to_string(message)
self.assertTrue("|type=03|" in message_str)
self.assertTrue("|body=Duck" in message_str)
@@ -53,4 +57,3 @@ class ServiceMessageTest(TestCase):
message = _f.of(ServiceMessageType.ROUTING_DATA, "Duck", "Donald")
serialized = _f.serialize(message)
assert RS + "0" + str(ServiceMessageType.ROUTING_DATA.value)
assert len(serialized) > 16
-251
View File
@@ -1,251 +0,0 @@
"""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()
-211
View File
@@ -1,211 +0,0 @@
import json
import unittest
from pydantic import BaseModel
from amqp.adapter.type_descriptor import TypeDescriptor, _instantiate
class TestClass(BaseModel):
value: str
extra: str
class TestTypeDescriptor(unittest.TestCase):
def test_annotated_with_json(self):
type_str = (
"typing.Optional[typing.Annotated[core.base.providers.ingestion.IngestionConfig, Json]]"
)
td = TypeDescriptor(type_str, {})
self.assertEqual(td.type, "core.base.providers.ingestion.IngestionConfig")
self.assertTrue(td.is_json)
def test_annotated_without_json(self):
type_str = "typing.Optional[typing.Annotated[str, SomethingElse]]"
td = TypeDescriptor(type_str, {})
self.assertEqual(td.type, "str")
self.assertFalse(td.is_json)
def test_annotated_multiple_args(self):
type_str = "typing.Optional[typing.Annotated[list[int], Json, more_args]]"
td = TypeDescriptor(type_str, {})
self.assertEqual(td.type, "list[int]")
self.assertTrue(td.is_json)
def test_simple_optional(self):
type_str = "typing.Optional[uuid.UUID]"
td = TypeDescriptor(type_str, {})
self.assertEqual(td.type, "uuid.UUID")
self.assertFalse(td.is_json)
def test_whitespace_variations(self):
test_cases = [
("typing.Optional[typing.Annotated[ str , Json ]]", "str", True),
("typing.Optional[ typing.Annotated[dict,Json]]", "dict", True),
("typing.Optional[ int ]", "int", False),
]
for type_str, expected_type, expected_json in test_cases:
with self.subTest(type_str=type_str):
td = TypeDescriptor(type_str, {})
self.assertEqual(td.type, expected_type)
self.assertEqual(td.is_json, expected_json)
def test_invalid_input(self):
test_cases = [
"typing.Optional[missing_bracket",
"typing.Optional[]",
"[without.optional]",
]
for type_str in test_cases:
with self.subTest(type_str=type_str):
td = TypeDescriptor(type_str, {})
self.assertEqual(td.type, "")
self.assertFalse(td.is_json)
def test_repr(self):
type_str = "typing.Optional[typing.Annotated[str, Json]]"
td = TypeDescriptor(type_str, {})
self.assertEqual(repr(td), "TypeDescriptor(type='str', is_json=True)")
def test_instantiate_basic_types(self):
test_cases = [
("str", "hello", "hello"),
("int", "42", 42),
("float", "3.14", 3.14),
("bool", "true", True),
]
for type_str, input_value, expected in test_cases:
with self.subTest(type_str=type_str, input_value=input_value):
result = _instantiate(type_str, input_value)
self.assertEqual(result, expected)
def test_instantiate_list_types(self):
test_cases = [
("typing.List[str]", "a,b,c", ["a", "b", "c"]),
("typing.List[int]", "1,2,3", [1, 2, 3]),
("typing.List[str]", json.dumps(["a", "b", "c"]), ["a", "b", "c"], True),
("typing.List[int]", json.dumps([1, 2, 3]), [1, 2, 3], True),
]
for type_str, input_value, expected, *args in test_cases:
is_json = args[0] if args else False
with self.subTest(type_str=type_str, input_value=input_value, is_json=is_json):
result = _instantiate(type_str, input_value, is_json=is_json)
self.assertEqual(result, expected)
def test_instantiate_dict_types(self):
test_cases = [
("typing.Dict", {"a": 1, "b": 2}, {"a": 1, "b": 2}),
("typing.Dict", json.dumps({"a": 1, "b": 2}), {"a": 1, "b": 2}, True),
("typing.Mapping", {"a": 1, "b": 2}, {"a": 1, "b": 2}),
]
for type_str, input_value, expected, *args in test_cases:
is_json = args[0] if args else False
with self.subTest(type_str=type_str, input_value=input_value, is_json=is_json):
result = _instantiate(type_str, input_value, is_json=is_json)
self.assertEqual(result, expected)
def test_instantiate_set_types(self):
test_cases = [
("typing.Set[str]", "a,b,c", {"a", "b", "c"}),
("typing.Set[int]", "1,2,3", {1, 2, 3}),
("typing.Set[str]", json.dumps(["a", "b", "c"]), {"a", "b", "c"}, True),
]
for type_str, input_value, expected, *args in test_cases:
is_json = args[0] if args else False
with self.subTest(type_str=type_str, input_value=input_value, is_json=is_json):
result = _instantiate(type_str, input_value, is_json=is_json)
self.assertEqual(result, expected)
def test_instantiate_tuple_types(self):
test_cases = [
("typing.Tuple[str, int]", ["hello", 42], ("hello", 42)),
("typing.Tuple[str, int]", json.dumps(["hello", 42]), ("hello", 42), True),
]
for type_str, input_value, expected, *args in test_cases:
is_json = args[0] if args else False
with self.subTest(type_str=type_str, input_value=input_value, is_json=is_json):
result = _instantiate(type_str, input_value, is_json=is_json)
self.assertEqual(result, expected)
def test_instantiate_optional_types(self):
test_cases = [
("typing.Optional[str]", "hello", "hello"),
("typing.Optional[typing.List[str]]", "a,b,c", ["a", "b", "c"]),
(
"typing.Optional[typing.List[str]]",
json.dumps(["a", "b", "c"]),
["a", "b", "c"],
True,
),
]
for type_str, input_value, expected, *args in test_cases:
is_json = args[0] if args else False
with self.subTest(type_str=type_str, input_value=input_value, is_json=is_json):
result = _instantiate(type_str, input_value, is_json=is_json)
self.assertEqual(result, expected)
def test_instantiate_nested_types(self):
test_cases = [
(
"typing.List[typing.Dict]",
json.dumps([{"a": 1}, {"b": 2}]),
[{"a": 1}, {"b": 2}],
True,
),
(
"typing.Optional[typing.List[typing.Dict]]",
json.dumps([{"a": 1}, {"b": 2}]),
[{"a": 1}, {"b": 2}],
True,
),
]
for type_str, input_value, expected, *args in test_cases:
is_json = args[0] if args else False
with self.subTest(type_str=type_str, input_value=input_value, is_json=is_json):
result = _instantiate(type_str, input_value, is_json=is_json)
self.assertEqual(result, expected)
def test_instantiate_error_handling(self):
test_cases = [
("typing.List[str]", "invalid json", True, json.JSONDecodeError),
("typing.NonExistentType", "value", False, ImportError),
]
for type_str, input_value, is_json, expected_exception in test_cases:
with self.subTest(type_str=type_str, input_value=input_value, is_json=is_json):
with self.assertRaises(expected_exception):
_instantiate(type_str, input_value, is_json=is_json)
def test_instantiate_with_extra_init_data(self):
test_cases = [
(
"tests.adapter.test_type_descriptor.TestClass",
{"value": "test"},
{"extra": "extra_value"},
lambda x: x.value == "test" and x.extra == "extra_value",
),
]
for type_str, input_value, extra_init_data, validator in test_cases:
with self.subTest(type_str=type_str, input_value=input_value):
result = _instantiate(type_str, input_value, extra_init_data=extra_init_data)
self.assertTrue(validator(result))
if __name__ == "__main__":
unittest.main()
-220
View File
@@ -1,220 +0,0 @@
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()
-209
View File
@@ -1,209 +0,0 @@
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")
-4
View File
@@ -1,4 +0,0 @@
pytest>=7.0.0
pytest-cov>=4.0.0
pytest-asyncio>=0.21.0
aio-pika>=9.0.0
+2 -2
View File
@@ -2,8 +2,8 @@ from unittest import TestCase
import pytest
from amqp.adapter.amq_route_factory import AMQRouteFactory
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.amq_route_factory import AMQRouteFactory
from amqp.model.model import AMQRoute
from amqp.router.route_database import RouteDatabase
@@ -80,7 +80,7 @@ class TestRouteDatabase(TestCase):
message = DataMessageFactory.get_instance(111).create_request_message(
"POST", "clever3-book.localhost", "/aa/bb?cc=d", {}, ""
)
route = self._database.find_route(message.domain(), message.path())
route = self._database.find_route(message.domain, message.path)
self.assertIsNotNone(route)
self.assertEqual(ROUTE_3, route.as_string())
route = self._database.find_route("clever3.book.localhost", "/three-hot-dogs")
-210
View File
@@ -1,210 +0,0 @@
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()
-164
View File
@@ -1,164 +0,0 @@
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)
-185
View File
@@ -1,185 +0,0 @@
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()
-12
View File
@@ -1,12 +0,0 @@
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
-90
View File
@@ -1,90 +0,0 @@
import asyncio
import aio_pika
import os
import argparse
import sys
from aio_pika.abc import AbstractRobustExchange
async def create_exchange(loop, rabbit_user, rabbit_password, rabbit_host):
"""
Asynchronously connects to RabbitMQ, declares a direct exchange,
and prints the result. Handles potential connection errors.
"""
exchange_name = "cleverthis.clevermicro.management"
connection = None # Keep track of the connection
try:
# Construct the connection URI from environment variables
connection_uri = f"amqp://{rabbit_user}:{rabbit_password}@{rabbit_host}/"
# Attempt to establish the connection
connection = await aio_pika.connect_robust(connection_uri, loop=loop)
print(f"Connected to RabbitMQ: {connection_uri}")
# Create a channel
channel = await connection.channel()
# Declare the exchange
exchange: AbstractRobustExchange = await channel.declare_exchange(
exchange_name, aio_pika.ExchangeType.DIRECT
)
print(f"Exchange '{exchange.name}' created successfully.")
except aio_pika.exceptions.AMQPConnectionError as e:
print(f"Error: Could not connect to RabbitMQ. Details: {e}")
sys.exit(1) # Exit with an error code
except aio_pika.exceptions.AMQPChannelError as e:
print(f"Error creating channel or exchange: {e}")
sys.exit(1)
except Exception as e:
print(f"An unexpected error occurred: {e}")
sys.exit(1)
finally:
# Ensure the connection is closed, even if errors occur
if connection:
await connection.close()
print("Connection to RabbitMQ closed.")
def main():
"""
Main function to parse command-line arguments, retrieve RabbitMQ
credentials from environment variables, and run the
async exchange creation.
"""
# Set up argument parsing
parser = argparse.ArgumentParser(description="Create a RabbitMQ exchange.")
parser.add_argument(
"--host",
default="127.0.0.1",
help="RabbitMQ host (default: localhost)",
)
# Parse the arguments
args = parser.parse_args()
rabbit_host = args.host
# Get RabbitMQ credentials from environment variables
rabbit_user = os.environ.get("RABBIT_MQ_USER")
rabbit_password = os.environ.get("RABBIT_MQ_PASSWORD")
if not rabbit_user or not rabbit_password:
print(
"Error: RabbitMQ credentials not found in environment variables.\n"
"Please set RABBIT_MQ_USER and RABBIT_MQ_PASSWORD."
)
sys.exit(1) # Exit with an error code
# Create the event loop
loop = asyncio.get_event_loop()
# Run the asynchronous exchange creation
loop.run_until_complete(create_exchange(loop, rabbit_user, rabbit_password, rabbit_host))
# Close the loop
loop.close()
if __name__ == "__main__":
main()