feat/#1 - applied 'black' formatting

This commit is contained in:
2025-04-18 23:49:10 +01:00
parent 04977e5ccc
commit 1de7afe4f5
35 changed files with 1239 additions and 563 deletions
+7 -4
View File
@@ -12,7 +12,7 @@ NULL_ROUTE = {
"queue": None,
"routingKey": "#",
"timeout": 1,
"validUntil": 0
"validUntil": 0,
}
@@ -49,18 +49,21 @@ class AMQRouteFactory:
queue=_queue,
key=_routing_key,
timeout=_timeout,
valid_until=_valid_until
valid_until=_valid_until,
)
logging_error(
"AMQRoute incorrect format. AMQRoute expects input as: "
"'Component-Name:AMQExchange-Name:Queue-Name:Routing-Key:Timeout:Valid-Until', "
"got: ['%s']", data
"got: ['%s']",
data,
)
except Exception as err:
logging_error(
"%s: AMQRoute incorrect format. AMQRoute expects input as: "
"'Component-Name:AMQExchange-Name:Queue-Name:Routing-Key:Timeout:Valid-Until', "
"got: ['%s']", err, data
"got: ['%s']",
err,
data,
)
return NULL_ROUTE
+155 -69
View File
@@ -2,6 +2,7 @@
The layer that interfaces with the CleverThis service that this Adapter is integrating with.
Methods here will invoke the actual CleverThis Service code, which are passed as Callable object.
"""
import asyncio
import inspect
import json
@@ -13,7 +14,12 @@ from typing import Dict, List, Tuple, Callable, Coroutine, Any
from urllib.parse import parse_qs, urlparse
from xml.etree import ElementTree
from aio_pika.abc import AbstractRobustChannel, AbstractRobustConnection, AbstractRobustExchange, AbstractExchange
from aio_pika.abc import (
AbstractRobustChannel,
AbstractRobustConnection,
AbstractRobustExchange,
AbstractExchange,
)
from aiohttp import ClientSession, ClientResponse
from dataclasses import dataclass
from fastapi import HTTPException
@@ -28,6 +34,7 @@ from amqp.adapter.file_uploader import publish_file_to_rabbitmq
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE
from IPython.core import debugger
debug = debugger.Pdb().set_trace
@@ -50,7 +57,9 @@ class AMQMessage(DataMessage):
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 {}
self.extra_data = (
data_message.extra_data if hasattr(data_message, "extra_data") else {}
)
def parse_get_url(url: str) -> Tuple[str, Dict[str, str]]:
@@ -68,7 +77,9 @@ def parse_get_url(url: str) -> Tuple[str, Dict[str, str]]:
_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
@@ -82,7 +93,7 @@ def parse_request_body(message: DataMessage):
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]
@@ -97,10 +108,13 @@ def parse_request_body(message: DataMessage):
# Parse URL-encoded form data
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()}
_form_data: dict = {
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'] and _form_data['form']:
if _form_data["files"] and _form_data["form"]:
_form_data = process_form_data(_form_data)
return _form_data
except Exception as e:
@@ -135,14 +149,15 @@ def parse_request_body(message: DataMessage):
async def _convert_file_response(
obj: FileResponse,
connection: AbstractRobustConnection,
loop: AbstractEventLoop | None
loop: AbstractEventLoop | None,
) -> 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.
"""
async def _wrap_amq_api_calls(
connection: AbstractRobustConnection
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
@@ -150,39 +165,53 @@ async def _convert_file_response(
_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 = 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}")
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
_future = asyncio.run_coroutine_threadsafe(_wrap_amq_api_calls(connection), loop)
while not _future.done():
await asyncio.sleep(0.1) # allow coroutine to execute. calling _future.result() will block the entire thread and the coroutine will never execute
await asyncio.sleep(
0.1
) # allow coroutine to execute. calling _future.result() will block the entire thread and the coroutine will never execute
_channel, _exchange, _routing_key = _future.result()
# publish_file also needs RabbitMQ API, ensure the loop is passed as argument
(_stream_name, _file_size) = await publish_file_to_rabbitmq(_exchange, obj.path, _routing_key, loop)
(_stream_name, _file_size) = await publish_file_to_rabbitmq(
_exchange, obj.path, _routing_key, loop
)
# as above, but luckily this time we don't need to wait for the result
_future = asyncio.run_coroutine_threadsafe(_channel.close(), loop)
return json.dumps({'files': [
{'filename': obj.filename,
'mediaType': 'application/octet-stream',
'streamName': _stream_name,
'fileSize': _file_size
}]
})
return json.dumps(
{
"files": [
{
"filename": obj.filename,
"mediaType": "application/octet-stream",
"streamName": _stream_name,
"fileSize": _file_size,
}
]
}
)
def is_likely_json(text: str) -> bool:
""" Quick test to see if should convert the response string to JSON format. """
"""Quick test to see if should convert the response string to JSON format."""
_text = text.strip()
return (_text.startswith('{') and _text.endswith('}')) or \
(_text.startswith('[') and _text.endswith(']'))
return (_text.startswith("{") and _text.endswith("}")) or (
_text.startswith("[") and _text.endswith("]")
)
async def call_cleverthis_api(
@@ -190,7 +219,7 @@ async def call_cleverthis_api(
args: Any,
api_method: Callable[[Any, Any], Coroutine],
success_code: int,
loop: AbstractEventLoop | None = None
loop: AbstractEventLoop | None = None,
) -> DataResponse:
"""
Handles POST and PUT requests, which typically contain possibly complex payload data, passed as a dictionary.
@@ -220,41 +249,69 @@ async def call_cleverthis_api(
else:
_verified_args[arg.name] = args[arg.name]
else:
logging.info(f"{api_method.__name__}: Missing required argument: {arg.name}")
logging.info(
f"{api_method.__name__}: Missing required argument: {arg.name}"
)
logging.info(
f" [*] INVOKE ENDPOINT: {api_method.__name__}, ARGS: {', '.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 = await _convert_file_response(
_result, message.connection, loop
), 'application/octet-stream'
_json_result, _content_type = (
await _convert_file_response(_result, message.connection, loop),
"application/octet-stream",
)
else:
_json_result = _result
_content_type = 'application/json'
_content_type = "application/json"
_binary_result = _json_result.encode('utf-8') if hasattr(
_json_result, 'encode'
) else _json_result.body if hasattr(
_json_result, 'body'
) else _json_result.read() if hasattr(_json_result, 'read') else _json_result
logging.info(f" [x] ALL DONE in call_cleverthis_api id:{message.id}, resp code:{success_code}")
return DataResponseFactory.of(message.id, success_code, _content_type, _binary_result, message.trace_info)
_binary_result = (
_json_result.encode("utf-8")
if hasattr(_json_result, "encode")
else (
_json_result.body
if hasattr(_json_result, "body")
else (
_json_result.read()
if hasattr(_json_result, "read")
else _json_result
)
)
)
logging.info(
f" [x] ALL DONE in call_cleverthis_api id:{message.id}, resp code:{success_code}"
)
return DataResponseFactory.of(
message.id, success_code, _content_type, _binary_result, message.trace_info
)
except HTTPException as http_error:
_content = str(http_error.detail)
if not is_likely_json(_content):
_content = '{"error": "' + _content + '"}'
logging_error(f"Service REST endpoint invocation failed: {http_error}")
return DataResponseFactory.of(message.id, http_error.status_code, "application/json",
_content.encode('utf-8'), message.trace_info)
return DataResponseFactory.of(
message.id,
http_error.status_code,
"application/json",
_content.encode("utf-8"),
message.trace_info,
)
except Exception as error:
logging_error(f"Service endpoint invocation failed: {error}")
_content = str(error)
if not is_likely_json(_content):
_content = '{"error": "' + _content + '"}'
return DataResponseFactory.of(message.id, 500, "application/json",
_content.encode('utf-8'), message.trace_info)
return DataResponseFactory.of(
message.id,
500,
"application/json",
_content.encode("utf-8"),
message.trace_info,
)
# ============= CleverThisServiceAdapter =============
@@ -266,7 +323,9 @@ class CleverThisServiceAdapter:
and override the methods to handle the specific API calls.
"""
def __init__(self, amq_configuration: AMQConfiguration, session: ClientSession = None):
def __init__(
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
@@ -277,8 +336,13 @@ class CleverThisServiceAdapter:
self.service_host = self.amq_configuration.dispatch.service_host
self.service_port = self.amq_configuration.dispatch.service_port
self.loop: AbstractEventLoop | None = None
self.require_authenticated_user = eval(
self.amq_configuration.amq_adapter.require_authenticated_user.capitalize()) if self.amq_configuration.amq_adapter.require_authenticated_user.strip() else True
self.require_authenticated_user = (
eval(
self.amq_configuration.amq_adapter.require_authenticated_user.capitalize()
)
if self.amq_configuration.amq_adapter.require_authenticated_user.strip()
else True
)
def get_content_type(self, message_headers: Dict[str, List[str]]) -> str:
"""
@@ -287,9 +351,9 @@ class CleverThisServiceAdapter:
return: The content type of the message.
"""
for header, values in message_headers.items():
if header.lower() == 'content-type':
return values[0].split(';')[0]
return 'application/json'
if header.lower() == "content-type":
return values[0].split(";")[0]
return "application/json"
async def get_current_user(self, token: str) -> object:
"""
@@ -310,37 +374,46 @@ class CleverThisServiceAdapter:
_amq_response: Any | None = None
logging.debug(
f' [*] on_message: _auth_user={_auth_user}, require_authenticated_user={self.require_authenticated_user}')
f" [*] on_message: _auth_user={_auth_user}, 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.info(f"Auth: {_auth}")
if isinstance(_auth, List):
_auth = _auth[0]
if 'Bearer' in _auth:
if "Bearer" in _auth:
try:
token = _auth.split(' ')[1]
logging.warn(f'Token: {token}')
token = _auth.split(" ")[1]
logging.warn(f"Token: {token}")
_auth_user = await self.get_current_user(token)
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'):
if (
_auth_user
or not self.require_authenticated_user
or message.path.endswith("/login")
):
method = message.method.lower()
if method == 'get':
if method == "get":
_amq_response = await self.get(message, _auth_user)
elif method == 'put':
elif method == "put":
_amq_response = await self.put(message, _auth_user)
elif method == 'post':
elif method == "post":
_amq_response = await self.post(message, _auth_user)
elif method == 'head':
elif method == "head":
_amq_response = await self.head(message, _auth_user)
elif method == 'delete':
elif method == "delete":
_amq_response = await self.delete(message, _auth_user)
else:
raise ValueError(f"Unexpected HTTP method: {message.method}")
else:
_amq_response = DataResponseFactory.of(message.id, 401, "text/plain", b"Unauthorized", message.trace_info)
logging.info(f" [x] ALL DONE in on_message id:{message.id}, resp code:{_amq_response.response_code}")
_amq_response = DataResponseFactory.of(
message.id, 401, "text/plain", b"Unauthorized", message.trace_info
)
logging.info(
f" [x] ALL DONE in on_message id:{message.id}, resp code:{_amq_response.response_code}"
)
return _amq_response
# ==================================================================================================
@@ -359,14 +432,25 @@ class CleverThisServiceAdapter:
# 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}")
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, _http_resp.status, _http_resp.content_type,
await _http_resp.read(),
message.trace_info)
return DataResponseFactory.of(message.id, 404, "text/plain",
str("Destination location does not exists").encode('utf-8'), message.trace_info)
return DataResponseFactory.of(
message.id,
_http_resp.status,
_http_resp.content_type,
await _http_resp.read(),
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: DataMessage, auth_user: Any) -> DataResponse:
"""
@@ -383,7 +467,7 @@ class CleverThisServiceAdapter:
data=message.body,
headers={**message.headers, **message.trace_info},
),
auth_user
auth_user,
)
async def post(self, message: DataMessage, auth_user: Any) -> DataResponse:
@@ -401,10 +485,12 @@ class CleverThisServiceAdapter:
data=message.body,
headers={**message.headers, **message.trace_info},
),
auth_user
auth_user,
)
async def handle_possible_form(self, message: DataMessage, request_coroutine, auth_user: Any) -> DataResponse:
async def handle_possible_form(
self, message: DataMessage, request_coroutine, auth_user: Any
) -> DataResponse:
_args = parse_request_body(message)
return await request_coroutine
+37 -19
View File
@@ -11,7 +11,12 @@ from typing import Dict, List
from amqp.model.model import DataMessage
from amqp.model.snowflake_id import SnowflakeId
from amqp.adapter.serializer import map_with_list_as_string, map_as_string, parse_map_list_string, parse_map_string
from amqp.adapter.serializer import (
map_with_list_as_string,
map_as_string,
parse_map_list_string,
parse_map_string,
)
from amqp.router.utils import sanitize_as_rabbitmq_name
from amqp.adapter.file_downloader import download_buffer
@@ -65,19 +70,24 @@ class DataMessageFactory:
self.snowflake_sequence = 0
self.snowflake_last_timestamp = timestamp
time = (timestamp - 1000 * int(DataMessageFactory.snowflake_epoch))
time = timestamp - 1000 * int(DataMessageFactory.snowflake_epoch)
return SnowflakeId(
time >> (64 - self.SEQUENCE_BITS - self.MACHINE_ID_BITS),
((time << (self.SEQUENCE_BITS + self.MACHINE_ID_BITS)) | self.machine_id | self.snowflake_sequence) & 0xFFFFFFFFFFFFFFFF
(
(time << (self.SEQUENCE_BITS + self.MACHINE_ID_BITS))
| self.machine_id
| self.snowflake_sequence
)
& 0xFFFFFFFFFFFFFFFF,
)
def create_request_message(
self,
method: str,
domain: str,
path: str,
headers: Dict[str, List[str]],
message,
self,
method: str,
domain: str,
path: str,
headers: Dict[str, List[str]],
message,
) -> DataMessage:
"""
Encapsulate / hide the system level inputs, create the AMQ message from business relevant input
@@ -89,7 +99,9 @@ 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(
self.MAGIC_BYTE_HTTP_REQUEST,
@@ -119,8 +131,12 @@ class DataMessageFactory:
:return: the message timestamp
"""
bits = _id.get_bits()
low = bits[0] >> (DataMessageFactory.SEQUENCE_BITS + DataMessageFactory.MACHINE_ID_BITS)
high = bits[1] << (64 - 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
)
return (high | low) + DataMessageFactory.snowflake_epoch * 1000
@staticmethod
@@ -152,13 +168,13 @@ class DataMessageFactory:
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="
).encode('utf-8')
).encode("utf-8")
header_length = 1 + len(id_bytes) + len(prolog)
prolog_len = header_length.to_bytes(2, byteorder='big')
prolog_len = header_length.to_bytes(2, byteorder="big")
msg_length = 2 + header_length + len(message.base64_body)
with BytesIO(bytearray(msg_length)) as baos:
baos.write(prolog_len)
baos.write(DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST.encode('utf-8'))
baos.write(DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST.encode("utf-8"))
baos.write(id_bytes)
baos.write(prolog)
baos.write(message.base64_body)
@@ -172,7 +188,7 @@ class DataMessageFactory:
:return: DataMessage instance
"""
prolog_len = (input_bytes[0] << 8) + input_bytes[1]
metadata = input_bytes[2:prolog_len + 2].decode()
metadata = input_bytes[2 : prolog_len + 2].decode()
pattern = re.compile(
f"{DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST}"
@@ -197,7 +213,7 @@ class DataMessageFactory:
path = match.group(4)
headers_str = match.group(5)
trace_info_str = match.group(6)
body_b64 = input_bytes[prolog_len + 2:]
body_b64 = input_bytes[prolog_len + 2 :]
headers = parse_map_list_string(headers_str)
trace_info = parse_map_string(trace_info_str)
@@ -223,9 +239,11 @@ class DataMessageFactory:
"""
# async implementation
message_data = json.loads(stream.decode())
req_info = message_data.get('req-info', '')
req_info = message_data.get("req-info", "")
if req_info:
(body, eof) = await download_buffer(loop=loop, rabbitmq_url=rabbitmq_url, queue_name=req_info)
(body, eof) = await download_buffer(
loop=loop, rabbitmq_url=rabbitmq_url, queue_name=req_info
)
if eof == 1:
return DataMessageFactory.from_bytes(body)
return None
+14 -8
View File
@@ -25,7 +25,13 @@ class DataResponseFactory:
)
@staticmethod
def of(id: SnowflakeId, response_code: int, content_type: str, body: bytes, trace_info: dict[str,str]) -> DataResponse:
def of(
id: SnowflakeId,
response_code: int,
content_type: str,
body: bytes,
trace_info: dict[str, str],
) -> DataResponse:
"""
Create the DataResponse message supplying all key values
:param id: SnowflakeID
@@ -42,7 +48,7 @@ class DataResponseFactory:
response=body,
error=None,
error_cause=None,
trace_info=trace_info
trace_info=trace_info,
)
@staticmethod
@@ -52,16 +58,16 @@ class DataResponseFactory:
:param response: Object to serialize
:return: byte array as serialized message
"""
id_bytes = response.id.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="
).encode('utf-8')
).encode("utf-8")
header_length = 1 + len(id_bytes) + len(prolog)
prolog_len = header_length.to_bytes(2, byteorder='big')
prolog_len = header_length.to_bytes(2, byteorder="big")
msg_length = 2 + header_length + len(response.response)
with BytesIO(bytearray(msg_length)) as baos:
baos.write(prolog_len)
baos.write(DataResponseFactory.MAGIC_BYTE_HTTP_REQUEST.encode('utf-8'))
baos.write(DataResponseFactory.MAGIC_BYTE_HTTP_REQUEST.encode("utf-8"))
baos.write(id_bytes)
baos.write(prolog)
baos.write(base64.b64encode(response.response))
@@ -75,7 +81,7 @@ class DataResponseFactory:
:return: DataResponse object
"""
prolog_len = (input_bytes[0] << 8) + input_bytes[1]
metadata = input_bytes[2:prolog_len + 2].decode()
metadata = input_bytes[2 : prolog_len + 2].decode()
pattern = re.compile(
f"{DataResponseFactory.MAGIC_BYTE_HTTP_REQUEST}"
@@ -100,7 +106,7 @@ class DataResponseFactory:
error = match.group(4)
error_cause = match.group(5)
trace_info_str = match.group(6)
body_b64 = input_bytes[prolog_len + 2:]
body_b64 = input_bytes[prolog_len + 2 :]
trace_info = parse_map_string(trace_info_str)
+58 -30
View File
@@ -52,7 +52,9 @@ def ensure_directory_exists(filepath) -> str:
version += 1
async def download_file(loop, rabbitmq_url, file_info, message_data, output_dir) -> (int, int):
async def download_file(
loop, rabbitmq_url, file_info, message_data, output_dir
) -> (int, int):
"""
Asynchronously downloads a file from RabbitMQ based on the provided file information.
takes the file_info and message_data and output_dir, creates queue consumer and waits for the file to be downloaded
@@ -64,17 +66,23 @@ async def download_file(loop, rabbitmq_url, file_info, message_data, output_dir)
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']
queue_name = file_info['queue_name']
filename = file_info["filename"]
size = file_info["size"]
stream_name = file_info["streamName"]
queue_name = file_info["queue_name"]
filepath = os.path.join(output_dir, filename)
filepath = ensure_directory_exists(filepath) # if filename exists, create a new (next) version
download_locations[file_info['key']] = filepath
filename = os.path.basename(filepath) # if a new version was created, ensure to use that new version
file_info['filename'] = filename
filepath = ensure_directory_exists(
filepath
) # if filename exists, create a new (next) version
download_locations[file_info["key"]] = filepath
filename = os.path.basename(
filepath
) # if a new version was created, ensure to use that new version
file_info["filename"] = filename
logging.info(f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name}:{queue_name} to {filepath}")
logging.info(
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)
@@ -87,10 +95,10 @@ async def download_file(loop, rabbitmq_url, file_info, message_data, output_dir)
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
with open(filepath, 'ab') as f:
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))
_eof = int(message.headers.get("eof", END_OF_FILE))
if _eof != IN_PROGRESS:
break
except Exception as e:
@@ -101,7 +109,9 @@ async def download_file(loop, rabbitmq_url, file_info, message_data, output_dir)
return _file_size, _eof
if queue_name:
_future: Future = asyncio.run_coroutine_threadsafe(wrap_rabbit_mq_calls(), loop=loop)
_future: Future = asyncio.run_coroutine_threadsafe(
wrap_rabbit_mq_calls(), loop=loop
)
while not _future.done():
await asyncio.sleep(0.05)
_local_file_size, _local_eof = _future.result()
@@ -109,6 +119,7 @@ async def download_file(loop, rabbitmq_url, file_info, message_data, output_dir)
return 0, CANCELLED
async def download_buffer(loop, rabbitmq_url, queue_name) -> (bytearray, int):
"""
Asynchronously downloads a bytearay from RabbitMQ based on the provided file information.
@@ -120,6 +131,7 @@ async def download_buffer(loop, rabbitmq_url, queue_name) -> (bytearray, int):
"""
logging.info(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
@@ -127,12 +139,14 @@ async def download_buffer(loop, rabbitmq_url, queue_name) -> (bytearray, int):
try:
channel: AbstractRobustChannel = await _connection.channel()
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
logging.info(f"Awaiting all buffer chunks being downloaded for {queue_name}")
logging.info(
f"Awaiting all buffer chunks being downloaded for {queue_name}"
)
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
_res.extend(message.body)
_eof = message.headers.get('eof', END_OF_FILE)
_eof = message.headers.get("eof", END_OF_FILE)
if _eof != IN_PROGRESS:
break
logging.info(f"Finished downloading buffer {queue_name}")
@@ -142,7 +156,9 @@ async def download_buffer(loop, rabbitmq_url, queue_name) -> (bytearray, int):
await _connection.close()
return _res, _eof
_future: Future = asyncio.run_coroutine_threadsafe(wrap_rabbit_mq_calls(), loop=loop)
_future: Future = asyncio.run_coroutine_threadsafe(
wrap_rabbit_mq_calls(), loop=loop
)
while not _future.done():
await asyncio.sleep(0.05)
_local_res, _local_eof = _future.result()
@@ -155,7 +171,7 @@ async def on_message(
json_data: bytes,
rabbitmq_url: str,
loop,
output_dir="downloaded_files"
output_dir="downloaded_files",
) -> defaultdict:
"""
Asynchronous callback function for handling messages from the 'amq-cleverswarm' queue.
@@ -170,35 +186,47 @@ async def on_message(
_amq_message_data = json.loads(json_data.decode())
files_to_download = []
for _key, _file_type in _amq_message_data['files'].items():
for _key, _file_type in _amq_message_data["files"].items():
if len(_file_type) > 0:
files_info = _file_type[0]
stream_name = files_info['streamName']
files_info['queue_name'] = _message_data.get(stream_name, "") if stream_name else ""
files_info['key'] = _key
stream_name = files_info["streamName"]
files_info["queue_name"] = (
_message_data.get(stream_name, "") if stream_name else ""
)
files_info["key"] = _key
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 False
# 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}")
_task = asyncio.create_task(download_file(loop, rabbitmq_url, file_info, _message_data, output_dir))
logging.info(
f" [{message.delivery_tag}] about to create task for {file_info}"
)
_task = asyncio.create_task(
download_file(loop, rabbitmq_url, file_info, _message_data, output_dir)
)
_tasks.append(_task)
logging.info(f" [{message.delivery_tag}] Waiting for all files to download...gather(*tasks)")
await asyncio.gather(*_tasks) # Wait for all downloads to complete
logging.info(
f" [{message.delivery_tag}] Waiting for all files to download...gather(*tasks)"
)
await asyncio.gather(*_tasks) # Wait for all downloads to complete
# After all files are downloaded, acknowledge the original message
logging.info(f"All files downloaded. d-tag: {message.delivery_tag}")
#logging.info(f"All files downloaded. Acknowledging the main message d-tag: {message.delivery_tag}")
#await message.ack()
#logging.info(f"Message d-tag: {message.delivery_tag} acknowledged.")
# logging.info(f"All files downloaded. Acknowledging the main message d-tag: {message.delivery_tag}")
# await message.ack()
# logging.info(f"Message d-tag: {message.delivery_tag} acknowledged.")
except Exception as e:
logging_error(f"Building DataMessage: {e}")
asyncio.run_coroutine_threadsafe(message.ack(), loop=loop) # ACK -> ignore the message (do not redeliver)
asyncio.run_coroutine_threadsafe(
message.ack(), loop=loop
) # ACK -> ignore the message (do not redeliver)
return download_locations
+27 -13
View File
@@ -20,7 +20,7 @@ async def publish_file_to_rabbitmq(
routing_key: str,
loop: AbstractEventLoop,
max_chunk_size: int = 2**20, # 1MB default chunk size
content_type: str = 'application/octet-stream', # Default content type
content_type: str = "application/octet-stream", # Default content type
delivery_mode: DeliveryMode = DeliveryMode.PERSISTENT,
) -> tuple[str, int]:
"""
@@ -45,7 +45,7 @@ async def publish_file_to_rabbitmq(
raise FileNotFoundError(f"File not found: {file_path}")
_file_size = os.path.getsize(file_path)
with open(file_path, 'rb') as f:
with open(file_path, "rb") as f:
_offset = 0
while True:
_chunk = f.read(max_chunk_size)
@@ -57,22 +57,36 @@ async def publish_file_to_rabbitmq(
content_type=content_type,
delivery_mode=delivery_mode,
headers={
'file_name': file_path.name,
'pos': _offset,
'total_size': _file_size,
'chunk_size': len(_chunk),
'eof': END_OF_FILE if max_chunk_size + _offset >= _file_size else IN_PROGRESS
}
"file_name": file_path.name,
"pos": _offset,
"total_size": _file_size,
"chunk_size": len(_chunk),
"eof": (
END_OF_FILE
if max_chunk_size + _offset >= _file_size
else IN_PROGRESS
),
},
)
logging.debug(
f"Going to Publish chunk {_offset}/{_file_size} to {exchange.name}:{routing_key}"
)
_future = asyncio.run_coroutine_threadsafe(
exchange.publish(message, routing_key=routing_key), loop
)
logging.debug(f"Going to Publish chunk {_offset}/{_file_size} to {exchange.name}:{routing_key}")
_future = asyncio.run_coroutine_threadsafe(exchange.publish(message, routing_key=routing_key), loop)
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(f"Published chunk {_offset}/{_file_size} to {exchange.name}:{routing_key}")
logging.debug(
f"Published chunk {_offset}/{_file_size} to {exchange.name}:{routing_key}"
)
logging_debug(f"File {file_path} published to RabbitMQ exchange {exchange.name}, queue {routing_key}")
logging_debug(
f"File {file_path} published to RabbitMQ exchange {exchange.name}, queue {routing_key}"
)
return routing_key, _file_size
+8 -2
View File
@@ -13,7 +13,10 @@ def logging_error(msg: str, *args) -> None:
_exec_info = sys.exc_info()[2]
if _exec_info:
_tb = traceback.extract_tb(sys.exc_info()[2])[-1]
logging.error(f" [E] '{msg}' ---> '{_tb.filename}': {_tb.lineno}, func: '{_tb.name}'", *args)
logging.error(
f" [E] '{msg}' ---> '{_tb.filename}': {_tb.lineno}, func: '{_tb.name}'",
*args,
)
else:
logging.error(f" [E] {msg}", *args)
@@ -28,6 +31,9 @@ def logging_debug(msg: str, *args) -> None:
_exec_info = sys.exc_info()[2]
if _exec_info:
_tb = traceback.extract_tb(sys.exc_info()[2])[-1]
logging.debug(f" [DBG] '{msg}' ---> '{_tb.filename}': {_tb.lineno}, func: '{_tb.name}'", *args)
logging.debug(
f" [DBG] '{msg}' ---> '{_tb.filename}': {_tb.lineno}, func: '{_tb.name}'",
*args,
)
else:
logging.debug(f" [DBG] {msg}", *args)
+29 -15
View File
@@ -33,12 +33,12 @@ def process_form_data(json_input) -> dict:
result = {}
# Process form data - take first element of each list
for key, value in data['form'].items():
for key, value in data["form"].items():
if value: # Only process if list is not empty
result[key] = value[0]
# Override with files data where available
for key, value in data['files'].items():
for key, value in data["files"].items():
if value: # Only override if files data is non-empty
result[key] = value[0]
@@ -52,32 +52,45 @@ class CleverMultiPartParser:
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()}
headers = {
key: "\n".join(value).encode("utf-8")
for key, value in message.headers.items()
}
content_type: str | bytes | None = headers.get("Content-Type")
content_type, params = multipart.parse_options_header(content_type)
if content_type.decode('utf-8') in ["application/octet-stream",
"multipart/form-data",
"application/x-www-form-urlencoded",
"application/x-url-encoded"]:
if content_type.decode("utf-8") in [
"application/octet-stream",
"multipart/form-data",
"application/x-www-form-urlencoded",
"application/x-url-encoded",
]:
try:
if hasattr(message, 'extra_data') and isinstance(message.extra_data, dict) and len(message.extra_data) > 0:
if (
hasattr(message, "extra_data")
and isinstance(message.extra_data, dict)
and len(message.extra_data) > 0
):
try:
self.form_data = json.loads(message.body())
self.is_json = True
if self.form_data['files'] and self.form_data['form']:
if self.form_data["files"] and self.form_data["form"]:
self.form_data = process_form_data(self.form_data)
except Exception as jsonerror:
logging_error(f"Parsing JSON: {jsonerror}")
self.is_valid = False
else:
python_multipart.parse_form(headers, io.BytesIO(message.body()), self.on_field, self.on_file)
python_multipart.parse_form(
headers, io.BytesIO(message.body()), self.on_field, self.on_file
)
self.is_multipart = True
except Exception as e:
logging_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'] and self.form_data['form']:
if self.form_data["files"] and self.form_data["form"]:
self.form_data = process_form_data(self.form_data)
except Exception as jsonerror:
logging_error(f"Parsing JSON: {jsonerror}")
@@ -87,11 +100,12 @@ class CleverMultiPartParser:
def on_field(self, field: Field):
logging_debug(field)
self.form_data[field.field_name.decode('utf-8')] = field.value
self.form_data[field.field_name.decode("utf-8")] = field.value
logging_debug(self.form_data)
def on_file(self, file: File):
logging_debug(file)
self.form_data[file.field_name.decode('utf-8')] = \
UploadFile(file.file_object, size=file.size, filename=file.file_name)
self.form_data[file.field_name.decode("utf-8")] = UploadFile(
file.file_object, size=file.size, filename=file.file_name
)
logging_debug(self.form_data)
+11 -6
View File
@@ -63,7 +63,7 @@ def deserialize_object(data: str) -> BaseModel:
An R2RSerializable object.
"""
class_name, dict_str = data.split(":", 1)
dict_str = dict_str.strip('{}')
dict_str = dict_str.strip("{}")
pairs = dict_str.split(",") if dict_str else []
obj_dict: Dict[str, Any] = {}
@@ -75,8 +75,8 @@ def deserialize_object(data: str) -> BaseModel:
obj_dict[k] = _convert_value(v)
logging_debug(f"Deserialized object: {class_name}")
if '[' in class_name:
class_name = class_name.split('[')[0]
if "[" in class_name:
class_name = class_name.split("[")[0]
logging_debug(f"trying object: {class_name}")
cls = globals().get(class_name)
if cls is None:
@@ -89,7 +89,7 @@ def _convert_value(v: str) -> Any:
"""
Helper function to convert a string to its appropriate Python type.
"""
if v == 'None':
if v == "None":
return None
elif _is_uuid_string(v):
return uuid.UUID(v)
@@ -97,8 +97,13 @@ def _convert_value(v: str) -> Any:
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
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)
+13 -13
View File
@@ -8,16 +8,16 @@ from typing import List, Union, Dict
def long_to_bytes(x: int) -> bytes:
return struct.pack('q', x)
return struct.pack("q", x)
def bytes_to_long(b: bytes) -> int:
return struct.unpack('q', b)[0]
return struct.unpack("q", b)[0]
def read_long(bis: BytesIO) -> int:
try:
return struct.unpack('q', bis.read(8))[0]
return struct.unpack("q", bis.read(8))[0]
except IOError:
return 0
@@ -29,37 +29,37 @@ def object_or_list_as_string(value: Union[str, List[str]]) -> str:
def map_as_string(map_: Dict[str, str]) -> str:
return ','.join(f"{k}={object_or_list_as_string(v)}" for k, v in map_.items())
return ",".join(f"{k}={object_or_list_as_string(v)}" for k, v in map_.items())
def map_with_list_as_string(map_: Dict[str, List[str]]) -> str:
return ','.join(f"{k}={object_or_list_as_string(v)}" for k, v in map_.items())
return ",".join(f"{k}={object_or_list_as_string(v)}" for k, v in map_.items())
def parse_map_string(input_str: str) -> Dict[str, str]:
map_ = {}
if input_str.startswith('{') and input_str.endswith('}'):
for entry in input_str[1:-1].split(','):
if input_str.startswith("{") and input_str.endswith("}"):
for entry in input_str[1:-1].split(","):
if len(entry) > 1:
key, value = entry.split('=')
key, value = entry.split("=")
map_[key] = value
return map_
def csv_as_list(input_str: str) -> List[str]:
return [item.strip() for item in input_str.strip('[]').split(',')]
return [item.strip() for item in input_str.strip("[]").split(",")]
def add_to_map(map_: Dict[str, List[str]], token: str) -> None:
eq_idx = token.index('=')
eq_idx = token.index("=")
if eq_idx > 0:
key = token[:eq_idx]
map_[key] = csv_as_list(token[eq_idx+1:])
map_[key] = csv_as_list(token[eq_idx + 1 :])
def parse_map_list_string(input_str: str) -> Dict[str, List[str]]:
map_ = {}
if input_str.startswith('{') and input_str.endswith('}'):
for token in input_str[1:-1].split('],'):
if input_str.startswith("{") and input_str.endswith("}"):
for token in input_str[1:-1].split("],"):
add_to_map(map_, token)
return map_
+27 -10
View File
@@ -20,6 +20,7 @@ class ServiceMessageFactory:
"""
build/serialize/deserialize CleverMicro Service Message class
"""
def __init__(self, name: str):
self.process_name = name
@@ -37,7 +38,9 @@ class ServiceMessageFactory:
)
return "null"
def of(self, message_type: ServiceMessageType, payload: str, recipient_name: str) -> ServiceMessage:
def of(
self, message_type: ServiceMessageType, payload: str, recipient_name: str
) -> ServiceMessage:
"""
Builds ServiceMessage object from relevant values.
:param message_type: enum Message type
@@ -53,10 +56,12 @@ class ServiceMessageFactory:
recipient_name=recipient_name,
reply_to=self.process_name,
trace_info=TraceInfoAdapter.create_produce_span(_id),
message=payload
message=payload,
)
def as_base64(self, message_type: ServiceMessageType, payload: str, recipient_name: str) -> ServiceMessage:
def as_base64(
self, message_type: ServiceMessageType, payload: str, recipient_name: str
) -> ServiceMessage:
"""
Builds ServiceMessage object from relevant values, stores content Base64 encoded.
:param message_type: enum Message type
@@ -72,7 +77,7 @@ class ServiceMessageFactory:
recipient_name=recipient_name,
reply_to=self.process_name,
trace_info=TraceInfoAdapter.create_produce_span(_id.as_string()),
message=payload
message=payload,
)
def from_bytes(self, input_bytes: bytes) -> ServiceMessage:
@@ -94,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(
@@ -104,14 +111,20 @@ 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(
"Unexpected ServiceMessage format. Expected 6 tokens separated by '|' in format: "
"[ID(64bit number)|MessageType(8bit number)|RecipientName(String)|"
"ReplyTo(String)|TraceInfo|Payload(String or Base64)], got: "
"[%s], reported problem: %s", input_bytes.decode(), str(e)
"[%s], reported problem: %s",
input_bytes.decode(),
str(e),
)
return INVALID_MESSAGE
@@ -125,7 +138,7 @@ class ServiceMessageFactory:
f"{message.id.as_string()}{RS}{self.format_message_type(message)}{RS}{message.recipient_name}{RS}"
f"{message.reply_to}{RS}{map_as_string(message.trace_info)}{RS}"
f"{base64.b64encode(message.message.encode()).decode() if message.payload_type != 0 else message.message}"
).encode('utf-8')
).encode("utf-8")
def format_message_type(self, message: ServiceMessage) -> str:
"""
@@ -134,6 +147,10 @@ class ServiceMessageFactory:
:param message: Service message to send
:return: formatted message type
"""
num_type = (message.message_type.value if message.message_type else ServiceMessageType.UNKNOWN.value) - 1
fmt = ((num_type & 0x7F) | (message.payload_type & 0x80))
num_type = (
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}"
+14 -4
View File
@@ -7,14 +7,24 @@ from amqp.model.model import AMQRoute
class TestAMQRouteFactory(TestCase):
def test_from_string(self):
_f = AMQRouteFactory()
_r: AMQRoute = _f.from_string("amq-clever-pybook::amq-clever-pybook:clever.pybook.localhost.#:5:0")
_r: AMQRoute = _f.from_string(
"amq-clever-pybook::amq-clever-pybook:clever.pybook.localhost.#:5:0"
)
self.assertEqual("", _r.exchange)
self.assertEqual("amq-clever-pybook", _r.queue)
self.assertEqual(5, _r.timeout)
_r2: AMQRoute = _f.from_string("amq-clever-pybook:amq-clever-pybook:clever.pybook.localhost.#:5")
_r2: AMQRoute = _f.from_string(
"amq-clever-pybook:amq-clever-pybook:clever.pybook.localhost.#:5"
)
self.assertEqual(NULL_ROUTE, _r2)
def test_validate(self):
_f = AMQRouteFactory()
self.assertTrue(_f.validate("amq-clever-pybook::amq-clever-pybook:clever.pybook.localhost.#:5:0"))
self.assertFalse(_f.validate("amq-clever-pybook:amq-clever-pybook:clever.pybook.localhost.#"))
self.assertTrue(
_f.validate(
"amq-clever-pybook::amq-clever-pybook:clever.pybook.localhost.#:5:0"
)
)
self.assertFalse(
_f.validate("amq-clever-pybook:amq-clever-pybook:clever.pybook.localhost.#")
)
+80 -19
View File
@@ -9,37 +9,76 @@ class DataMessageFactoryTest(TestCase):
raw = "\002\027A64f3b5af4d00000a4000|m=POST|d=clever-pybook.localhost|p=/debug|h={Accept=[*/*],X-Forwarded-Host=[clever-pybook.localhost:8085],User-Agent=[curl/8.7.1],X-Forwarded-Proto=[http],X-Forwarded-For=[172.18.0.1],Host=[clever-pybook.localhost:8085],Accept-Encoding=[gzip],Content-Length=[1475],X-Forwarded-Port=[8085],X-Real-Ip=[172.18.0.1],X-Forwarded-Server=[c5945bef67d1],Content-Type=[multipart/form-data; boundary=------------------------j1RjIqLrFH07XMsyMvpS1v]}|t={traceparent=00-bed6eb7de0d8ccf4c80b145a4bfec357-8d1f3bafc471bf0f-01}|b=TEST_ME"
msg = DataMessageFactory.from_bytes(raw.encode("utf-8"))
self.assertIsNotNone(msg.id)
self.assertEqual(msg.id.as_string(),'64F3B5AF4D00000A4000'.lower(),'SnowflakeID incorrectly parsed.')
self.assertEqual(msg.headers["Content-Type"][0], "multipart/form-data; boundary=------------------------j1RjIqLrFH07XMsyMvpS1v")
self.assertEqual(
msg.id.as_string(),
"64F3B5AF4D00000A4000".lower(),
"SnowflakeID incorrectly parsed.",
)
self.assertEqual(
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()
_i1 = _f.generate_snowflake_id()
_i2 = _f.generate_snowflake_id()
_i3 = _f.generate_snowflake_id()
self.assertEqual('0007F000'.lower(), _i0.as_string()[-8:], 'Instance ID and sequence not set correctly')
self.assertEqual('0007F001'.lower(), _i1.as_string()[-8:], 'Instance ID and sequence not set correctly')
self.assertEqual('0007F002'.lower(), _i2.as_string()[-8:], 'Instance ID and sequence not set correctly')
self.assertEqual('0007F003'.lower(), _i3.as_string()[-8:], 'Instance ID and sequence not set correctly')
self.assertEqual(
"0007F000".lower(),
_i0.as_string()[-8:],
"Instance ID and sequence not set correctly",
)
self.assertEqual(
"0007F001".lower(),
_i1.as_string()[-8:],
"Instance ID and sequence not set correctly",
)
self.assertEqual(
"0007F002".lower(),
_i2.as_string()[-8:],
"Instance ID and sequence not set correctly",
)
self.assertEqual(
"0007F003".lower(),
_i3.as_string()[-8:],
"Instance ID and sequence not set correctly",
)
def test_create_request_message(self):
_f = DataMessageFactory(34)
_req = _f.create_request_message("PUT", "localhost", "/test/me", {"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(b'Quick brown fox', _req.body(), "Body not correctly decoded")
_req = _f.create_request_message(
"PUT",
"localhost",
"/test/me",
{"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(b"Quick brown fox", _req.body(), "Body not correctly decoded")
def test_serialize(self):
_f = DataMessageFactory(34)
_req = _f.create_request_message("PUT", "localhost", "/test/me", {"k1": ["v1"], "k2": ["v2"]}, "Quick brown fox")
_req = _f.create_request_message(
"PUT",
"localhost",
"/test/me",
{"k1": ["v1"], "k2": ["v2"]},
"Quick brown fox",
)
_ser = _f.serialize(_req)
self.assertEqual(65, _ser[2])
_deser = _f.from_bytes(_ser)
@@ -51,18 +90,40 @@ class DataMessageFactoryTest(TestCase):
def test_amq_message_routing_key(self):
_f = DataMessageFactory(35)
_req = _f.create_request_message("PUT", "localhost", "/test/me", {"k1": ["v1"], "k2": ["v2"]}, "Quick brown fox")
_req = _f.create_request_message(
"PUT",
"localhost",
"/test/me",
{"k1": ["v1"], "k2": ["v2"]},
"Quick brown fox",
)
_key = _f.amq_message_routing_key(_req)
self.assertEqual("localhost.test.me", _key, "Message routing key wrongly extracted from the message")
self.assertEqual(
"localhost.test.me",
_key,
"Message routing key wrongly extracted from the message",
)
def test_get_timestamp_from_id(self):
_f = DataMessageFactory(35)
_timestamp = _f.get_timestamp_from_id(_f.generate_snowflake_id())
_x_timestamp = int(datetime.now(timezone.utc).timestamp() * 1000)
self.assertLessEqual(_x_timestamp - _timestamp, 1, "Timestamp in SnowflakeId differs too much from current time")
self.assertLessEqual(
_x_timestamp - _timestamp,
1,
"Timestamp in SnowflakeId differs too much from current time",
)
def test_to_string(self):
_f = DataMessageFactory(35)
_req = _f.create_request_message("PUT", "localhost", "/test/me", {"k1": ["v1"], "k2": ["v2"]}, "Quick brown fox")
_req = _f.create_request_message(
"PUT",
"localhost",
"/test/me",
{"k1": ["v1"], "k2": ["v2"]},
"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"
)
+7 -3
View File
@@ -20,9 +20,13 @@ class TestDataResponseFactory(TestCase):
_snowflakeId: SnowflakeId = _g.generate_snowflake_id()
_resp = _f.create_async_response_message(_snowflakeId)
_ser = _f.serialize(_resp)
self.assertEqual(2, str(_ser[2:]).index(_f.MAGIC_BYTE_HTTP_REQUEST+_snowflakeId.as_string()), "Magic byte and SnowflakeId not found")
self.assertEqual(
2,
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)
self.assertEqual((_ser[0]+_ser[1]), len(_ser) - 2 - len('T0s='))
self.assertEqual((_ser[0] + _ser[1]), len(_ser) - 2 - len("T0s="))
def test_from_bytes(self):
_f = DataResponseFactory()
@@ -33,4 +37,4 @@ class TestDataResponseFactory(TestCase):
_deser = _f.from_bytes(_ser)
self.assertEqual(_snowflakeId.as_string(), _deser.id)
self.assertEqual(200, _deser.response_code)
self.assertEqual(b'OK', _deser.response)
self.assertEqual(b"OK", _deser.response)
+7 -2
View File
@@ -6,6 +6,7 @@ import pytest
from amqp.adapter.file_downloader import ensure_directory_exists
class TestFileDownloader(TestCase):
def test_ensure_directory_exists(self):
"""
@@ -30,7 +31,9 @@ class TestFileDownloader(TestCase):
with open(filepath2, "w") as f:
f.write("test")
result2_v1 = ensure_directory_exists(filepath2)
assert result2_v1 == os.path.join(temp_dir, "test2.1.txt"), f"Test Case 2 Version 1 Failed: {result2_v1} != {os.path.join(temp_dir, 'test2.1.txt')}"
assert result2_v1 == os.path.join(
temp_dir, "test2.1.txt"
), f"Test Case 2 Version 1 Failed: {result2_v1} != {os.path.join(temp_dir, 'test2.1.txt')}"
# Create a few versioned files
with open(os.path.join(temp_dir, "test3.txt"), "w") as f:
@@ -41,6 +44,8 @@ class TestFileDownloader(TestCase):
f.write("test")
filepath3 = os.path.join(temp_dir, "test3.txt")
result3_v3 = ensure_directory_exists(filepath3)
assert result3_v3 == os.path.join(temp_dir, "test3.3.txt"), f"Test Case 3 Version 3 Failed: {result3_v3} != {os.path.join(temp_dir, 'test3.3.txt')}"
assert result3_v3 == os.path.join(
temp_dir, "test3.3.txt"
), f"Test Case 3 Version 3 Failed: {result3_v3} != {os.path.join(temp_dir, 'test3.3.txt')}"
print("All test cases passed!")
+7 -3
View File
@@ -28,9 +28,13 @@ class PydanticSerializerTest(TestCase):
assert False
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")
path, args = parse_url_query(
"http://localhost:8080/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741?response_type=JsonFile"
)
assert path == "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741"
assert args == {'response_type': 'JsonFile'}
path, args = parse_url_query("http://localhost:8080/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"
)
assert path == "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741"
assert args == {}
+15 -7
View File
@@ -1,6 +1,10 @@
from unittest import TestCase
from amqp.adapter.service_message_factory import ServiceMessageFactory, RS, INVALID_MESSAGE
from amqp.adapter.service_message_factory import (
ServiceMessageFactory,
RS,
INVALID_MESSAGE,
)
from amqp.model.model import ServiceMessage
from amqp.model.service_message_type import ServiceMessageType
@@ -8,17 +12,21 @@ from amqp.model.service_message_type import ServiceMessageType
class ServiceMessageTest(TestCase):
def test_to_string(self):
_f: ServiceMessageFactory = ServiceMessageFactory('bumble-bee')
message: ServiceMessage = _f.of(ServiceMessageType.ROUTING_DATA, "Duck", "Donald")
_f: ServiceMessageFactory = ServiceMessageFactory("bumble-bee")
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)
self.assertTrue("|replyTo=bumble-bee" in message_str)
nullMsg = _f.to_string(None)
nullMsg = _f.to_string(None)
self.assertEqual("null", nullMsg)
def test_from(self):
@@ -31,7 +39,7 @@ class ServiceMessageTest(TestCase):
assert deserialized.recipient_name == "Donald"
assert deserialized.message == "Duck"
assert deserialized.reply_to == "amqp-router"
invalidMsg = _f.from_bytes(b'wrong')
invalidMsg = _f.from_bytes(b"wrong")
self.assertEqual(invalidMsg, INVALID_MESSAGE)
def test_base64(self):
@@ -43,7 +51,7 @@ class ServiceMessageTest(TestCase):
assert deserialized.is_valid()
assert deserialized.message_type == ServiceMessageType.ROUTING_DATA
assert deserialized.message == "Duck"
def test_serialize(self):
_f: ServiceMessageFactory = ServiceMessageFactory("amqp-router")
message = _f.of(ServiceMessageType.ROUTING_DATA, "Duck", "Donald")
+2 -3
View File
@@ -1,6 +1,7 @@
from dataclasses import dataclass
from typing import Dict
@dataclass
class TraceInfoAdapter:
"""
@@ -9,6 +10,4 @@ class TraceInfoAdapter:
@staticmethod
def create_produce_span(message_id: str) -> Dict[str, str]:
return {
"trace-id": message_id
}
return {"trace-id": message_id}