Compare commits

...

44 Commits

Author SHA1 Message Date
hurui200320 45941de6b1 fix(General): fix availability report
/ build-and-push (push) Successful in 1m21s
Unit test coverage / pytest (push) Failing after 1m27s
Unit test coverage / pytest (pull_request) Failing after 1m17s
Remove the loop for monitoring the CPU usage, adopt changes in the management service for the

reporting message format, update the timing for sending out report, similar to java adapter.

ISSUES RELATED: clevermicro/amq-adapter-python#58
2025-07-15 15:07:52 +08:00
Stanislav Hejny 5075f510e9 add consul configuration
Unit test coverage / pytest (push) Failing after 1m19s
/ build-and-push (push) Successful in 1m25s
Unit test coverage / pytest (pull_request) Failing after 1m24s
Build and Publish Docker Image / build-and-push (push) Has been cancelled
2025-07-11 22:41:40 +01:00
Stanislav Hejny 99b7f73447 refactor: unify config retrieval in Consul ID generator functions
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-11 22:39:43 +01:00
Stanislav Hejny 03983306b1 refactor: use AMQConfiguration for Consul settings in ID generators
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-11 22:39:43 +01:00
Stanislav Hejny d0b580b136 feat: add Consul-based global ID generator with fallback mechanism
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-11 22:39:43 +01:00
Stanislav Hejny 287395a015 feat: add Consul-based global ID generator with fallback mechanism 2025-07-11 22:39:43 +01:00
hurui200320 38e5b532cf fix: update routing key
Unit test coverage / pytest (push) Failing after 1m14s
/ build-and-push (push) Successful in 1m24s
Build and Publish Docker Image / build-and-push (push) Has been cancelled
2025-07-09 15:25:04 +08:00
hurui200320 771154facd fix(General): fix for cm-dev deployment
Build and Publish Docker Image / build-and-push (push) Waiting to run
Unit test coverage / pytest (push) Failing after 1m32s
/ build-and-push (push) Successful in 1m35s
2025-07-08 18:38:40 +08:00
Stanislav Hejny fde6f17b58 Use image from repo rather than local build
Unit test coverage / pytest (push) Failing after 1m15s
/ build-and-push (push) Successful in 1m26s
Build and Publish Docker Image / build-and-push (push) Has been cancelled
2025-07-04 10:05:41 +01:00
Stanislav Hejny ec7cf6b924 fix: correct package name to python-multipart in pyproject.toml
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-04 10:05:41 +01:00
Stanislav Hejny 783c8a42d9 build: optimize Dockerfile with multi-stage build
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-04 10:05:41 +01:00
hurui200320 ed75f5d002 upgrade python_multipart version for demo
Build and Publish Docker Image / build-and-push (push) Waiting to run
Unit test coverage / pytest (push) Failing after 1m13s
/ build-and-push (push) Successful in 1m28s
2025-07-04 17:00:29 +08:00
Stanislav Hejny 2ab2509a55 feat #58 - publish docker image with demo app, fix test case
Unit test coverage / pytest (push) Failing after 1m19s
/ build-and-push (push) Successful in 1m26s
Build and Publish Docker Image / build-and-push (push) Has been cancelled
2025-07-01 22:48:22 +01:00
Stanislav Hejny 98129b86d0 feat #58 - publish docker image with demo app
Build and Publish Docker Image / build-and-push (push) Waiting to run
Unit test coverage / pytest (push) Failing after 1m20s
/ build-and-push (push) Successful in 1m30s
2025-07-01 21:45:11 +01:00
Stanislav Hejny 00018c79ac feat: add CI/CD, Docker Compose, and update demo for Docker environment
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-01 20:29:50 +01:00
Stanislav Hejny 7fd1c62596 Feat #58 - fix backpressure logic and update tests
Unit test coverage / pytest (push) Failing after 1m17s
/ build-and-push (push) Successful in 1m29s
2025-06-28 10:20:55 +02:00
stanislav.hejny a2d76aca2f Merge pull request 'otel-demo' (#56) from otel-demo into develop
Unit test coverage / pytest (push) Successful in 1m39s
CI for pypl publish / publish-lib (push) Failing after 1m45s
/ build-and-push (push) Successful in 1m51s
Reviewed-on: #56
2025-06-20 14:19:15 +00:00
Stanislav Hejny f13fd4153a Fix dependency versions
Unit test coverage / pytest (push) Successful in 1m21s
/ build-and-push (push) Successful in 1m45s
Unit test coverage / pytest (pull_request) Successful in 1m21s
2025-06-20 16:17:12 +02:00
Stanislav Hejny e280a3e13a add publish docker image action
/ build-and-push (push) Failing after 1m5s
Unit test coverage / pytest (push) Successful in 1m16s
2025-06-20 15:49:09 +02:00
stanislav.hejny b5918dcf45 Merge pull request 'try publish demo app as artefact' (#55) from otel-demo into develop
Unit test coverage / pytest (push) Successful in 1m25s
CI for pypl publish / publish-lib (push) Failing after 1m36s
Reviewed-on: #55
2025-06-20 13:33:52 +00:00
Stanislav Hejny e180cbe2df try publish demo app as artefact
Unit test coverage / pytest (push) Successful in 1m17s
Unit test coverage / pytest (pull_request) Successful in 1m26s
2025-06-20 15:22:30 +02:00
Stanislav Hejny 83edca0624 Merge branch 'cleverbrag-integration-2' into develop
Unit test coverage / pytest (push) Successful in 1m16s
CI for pypl publish / publish-lib (push) Successful in 1m19s
2025-06-03 00:33:16 +01:00
stanislav.hejny bfa6a97ca8 Merge branch 'develop' into cleverbrag-integration-2
Unit test coverage / pytest (push) Successful in 1m14s
Unit test coverage / pytest (pull_request) Successful in 1m14s
2025-06-02 23:23:26 +00:00
Stanislav Hejny 748af5a3d7 Issue #23 - add test suite for CleverBRAG
Unit test coverage / pytest (push) Successful in 1m6s
Unit test coverage / pytest (pull_request) Successful in 1m11s
2025-06-02 12:18:14 +01:00
Stanislav Hejny 7e04106ab4 Issue #23 - add unit test for new code from fixes identified after running test for BRAG collections
Unit test coverage / pytest (push) Successful in 1m8s
2025-06-02 12:11:07 +01:00
Stanislav Hejny d9583cb218 Issue #23 - fixes identified after running test for BRAG collections 2025-06-02 11:09:10 +01:00
Stanislav Hejny 8db2ad0b64 issue #23 - unit test coverage
Unit test coverage / pytest (push) Failing after 3m17s
2025-05-29 01:43:21 +01:00
Stanislav Hejny 5261b1f959 issue #23 - fixing issues from integration test, part 3 2025-05-29 01:05:05 +01:00
Stanislav Hejny 59e04300b0 issue #23 - fixing issues from integration test, part 2 2025-05-28 11:14:10 +01:00
Stanislav Hejny 82c8a298cd issue #23 - fixing issues from integration test
Unit test coverage / pytest (push) Failing after 1m9s
2025-05-27 13:34:09 +01:00
hurui200320 8224e4020f ci: move back to forgejo
Unit test coverage / pytest (pull_request) Successful in 2m3s
Unit test coverage / pytest (push) Successful in 1m51s
CI for pypl publish / publish-lib (push) Successful in 1m53s
ISSUES CLOSED: clevermicro/amq-adapter-python#8
2025-05-22 04:42:03 +00:00
stanislav.hejny 73aca8cf25 Merge pull request 'Implement unit test' (#47) from unittest-coverage-85 into develop
Unit test coverage / pytest (push) Successful in 1m14s
CI for pypl publish / publish-lib (push) Failing after 1m15s
Reviewed-on: #47
2025-05-19 14:30:14 +00:00
Stanislav Hejny 3e0f769837 feat #16 - update version so the code can be published as artifact
Unit test coverage / pytest (push) Successful in 1m19s
Unit test coverage / pytest (pull_request) Successful in 1m17s
2025-05-19 15:27:37 +01:00
hurui200320 2e12d5ab57 test(General): finished unit test, coverage 88%
Unit test coverage / pytest (push) Successful in 1m8s
Unit test coverage / pytest (pull_request) Successful in 1m6s
ISSUES CLOSED: clevermicro/amq-adapter-python#16
2025-05-16 18:16:04 +08:00
hurui200320 aec53f5b34 test(General): implement unit test for data parser, rabbitmq client and other python files
Unit test coverage / pytest (push) Failing after 1m7s
This commit bump the coverage rate to 77%
2025-05-15 17:06:05 +08:00
Stanislav Hejny 33431e87fc issue #16 - add model tests, cursor config file and update the tracing.py to setup Jaeger & Prometheus exporters
Unit test coverage / pytest (push) Failing after 1m4s
2025-05-15 00:47:12 +01:00
hurui200320 883ea3400f test(General): fix existing unit tests for cleverthis service adapter
Unit test coverage / pytest (push) Failing after 1m10s
Should bump coverage rate to 62%
2025-05-14 16:16:10 +08:00
hurui200320 85297fa4cf test(General): finished unit test for file handler, fixed unit test for backpressure handler
Unit test coverage / pytest (push) Failing after 1m4s
2025-05-14 15:08:32 +08:00
Stanislav Hejny 53dce69e8f issue #16 - fix the pytest hanging after test main thread exits
Unit test coverage / pytest (push) Failing after 1m4s
2025-05-14 01:01:35 +01:00
hurui200320 98fc52fc26 Partially finished the unit test for file handler
Unit test coverage / pytest (push) Failing after 3h4m2s
2025-05-13 18:00:01 +08:00
hurui200320 44b6a8c3f1 test(General): implement unit test for serializer 2025-05-13 15:14:50 +08:00
hurui200320 45357c42a8 test(General): write test for pydantic serializer, fixed serialization bug
Unit test coverage / pytest (push) Failing after 3h11m8s
This commit write unit tests for pydantic serializer using cursor ai, and found an issue related to serialization. Originally, the code uses comma to concat string, but map and list also has comma inside. During the deserialization, split using comma will break map and list, leaving `{` with first kv pair for map, and `[` with first element with list. Instead of using concat with comma as separator, I use json to encode the object dict.
2025-05-12 16:42:34 +08:00
Stanislav Hejny bc98351493 feat \#16 - add unittest for backpressure_handler and logging_utils
Unit test coverage / pytest (push) Failing after 3h11m10s
2025-05-12 01:06:39 +01:00
stanislav.hejny 5684deb65e Merge pull request 'WIP #23 - first draft of CleverBRAG facade with endpoints discovery, suggested type (de)serialization and sample invocation' (#24) from cleverbrag-integration into develop
Unit test coverage / pytest (push) Failing after 1m49s
CI for pypl publish / publish-lib (push) Failing after 1m35s
Reviewed-on: #24
2025-05-11 23:47:46 +00:00
66 changed files with 6732 additions and 593 deletions
+118
View File
@@ -0,0 +1,118 @@
---
description:
globs:
alwaysApply: true
---
---
description: Python best practices and patterns for modern software development with Flask and SQLite
globs: **/*.py, src/**/*.py, tests/**/*.py
---
# Python Best Practices
## Project Structure
- Use src-layout with `amqp/`
- Place tests in `tests/` directory parallel to `amqp/`
- Keep configuration in `amqp/config/` or as environment variables
- Store requirements in `pyproject.toml`
- Place static files in `static/` directory
- Use `templates/` for Jinja2 templates
## Code Style
- Follow Black code formatting
- Use isort for import sorting
- Follow PEP 8 naming conventions:
- snake_case for functions and variables
- PascalCase for classes
- UPPER_CASE for constants
- Maximum line length of 100 characters
- Use absolute imports over relative imports
## Type Hints
- Use type hints for all function parameters and returns
- Import types from `typing` module
- Use `Optional[Type]` instead of `Type | None`
- Use `TypeVar` for generic types
- Define custom types in `types.py`
- Use `Protocol` for duck typing
## Database
- Use SQLAlchemy ORM
- Implement database migrations with Alembic
- Use proper connection pooling
- Define models in separate modules
- Implement proper relationships
- Use proper indexing strategies
## Authentication
##- Use Flask-Login for session management
##- Implement Google OAuth using Flask-OAuth
##- Hash passwords with bcrypt
##- Use proper session security
##- Implement CSRF protection
##- Use proper role-based access control
## API Design
##- Use Flask-RESTful for REST APIs
##- Implement proper request validation
##- Use proper HTTP status codes
##- Handle errors consistently
##- Use proper response formats
##- Implement proper rate limiting
## Testing
- Use pytest for testing
- Write tests for all routes
- Use pytest-cov for coverage
- Implement proper fixtures
- Use proper mocking with pytest-mock
- Test all error scenarios
## Security
##- Use HTTPS in production
##- Implement proper CORS
##- Sanitize all user inputs
##- Use proper session configuration
- Implement proper logging
##- Follow OWASP guidelines
## Performance
##- Use proper caching with Flask-Caching
##- Implement database query optimization
##- Use proper connection pooling
##- Implement proper pagination
##- Use background tasks for heavy operations
- Monitor application performance
## Error Handling
- Create custom exception classes
- Use proper try-except blocks
- Implement proper logging
- Return proper error responses
- Handle edge cases properly
- Use proper error messages
## Documentation
- Use Google-style docstrings
- Document all public APIs
- Keep README.md updated
- Use proper inline comments
- Generate API documentation
- Document environment setup
## Development Workflow
- Use virtual environments (venv)
- Use or Implement pre-commit hooks
- Use proper Git workflow
- Follow semantic versioning
- Use proper CI/CD practices
- Implement proper logging
## Dependencies
##- Pin dependency versions
##- Use requirements.txt for production
- Separate dev dependencies
##- Use proper package versions
##- Regularly update dependencies
- Check for security vulnerabilitie
]cvx -=
+26
View File
@@ -0,0 +1,26 @@
.git
.github
.forgejo
__pycache__
*.pyc
*.pyo
*.pyd
.Python
env/
venv/
.venv/
pip-log.txt
pip-delete-this-directory.txt
.tox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.log
.git*
tests/
container-data/
docker-compose.yml
.dockerignore
+41
View File
@@ -0,0 +1,41 @@
on: [push]
env:
REGISTRY_URL: "git.cleverthis.com"
REPOSITORY: "clevermicro/amq-adapter-python-demo"
DOCKER_HOST: "tcp://dind:2375"
jobs:
build-and-push:
runs-on: general
services:
dind:
image: docker:dind
cmd:
- dockerd
- -H
- tcp://0.0.0.0:2375
- --tls=false
container:
image: "ghcr.io/catthehacker/ubuntu:js-22.04"
steps:
- name: checkout
uses: actions/checkout@v4
with:
ref: ${{ github.ref_name }}
- name: docker login
uses: docker/login-action@v3
with:
registry: ${{ env.REGISTRY_URL }}
username: ${{ secrets.REGISTRY_USER }}
password: ${{ secrets.REGISTRY_PASSWORD }}
- name: docker build & push
uses: docker/build-push-action@v6
with:
context: .
file: ./Dockerfile
no-cache: true
push: true
tags: ${{ env.REGISTRY_URL }}/${{ env.REPOSITORY }}:20250715-01
+3 -3
View File
@@ -40,9 +40,9 @@ jobs:
# publish
- name: publish to pulp pypi
run: |
twine upload --repository-url https://pkg.cleverthis.io/pypi/clevermicro/amq-adapter-python/simple/ \
--username $CI_REGISTRY_USER \
--password $CI_REGISTRY_PASSWORD \
twine upload --repository-url https://git.cleverthis.com/api/packages/clevermicro/pypi \
--username ${CI_REGISTRY_USER} \
--password ${CI_REGISTRY_PASSWORD} \
--verbose \
dist/*
env:
+2 -1
View File
@@ -56,7 +56,8 @@ amq_adapter.egg-info/
build/
unused-code/
.vscode/
amq_adapter.egg-info/
build/
coverage.xml
.aider*
+49
View File
@@ -0,0 +1,49 @@
# Build stage
FROM python:3.11-slim AS builder
WORKDIR /app
ENV PYTHONPATH=/app
# Install system dependencies for building
RUN apt-get update && apt-get install -y gcc && rm -rf /var/lib/apt/lists/*
# Copy project files for building
COPY pyproject.toml .
COPY . .
RUN pip install --no-cache-dir build && python -m build --wheel
# Runtime stage
FROM python:3.11-slim
WORKDIR /app
ENV PYTHONPATH=/app
RUN apt-get update && apt-get install -y libpq-dev && rm -rf /var/lib/apt/lists/*
# Install the built package from builder
COPY --from=builder /app/dist/*.whl /tmp/
RUN pip install --no-cache-dir /tmp/*.whl
# Install additional runtime dependencies for demo.py
RUN pip install --no-cache-dir \
uvicorn[standard]==0.29.0 \
python-multipart==0.0.20 \
httpx==0.27.0 \
opentelemetry-api==1.34.1 \
opentelemetry-sdk==1.34.1 \
opentelemetry-exporter-otlp-proto-grpc==1.34.1 \
opentelemetry-exporter-prometheus==0.55b1 \
opentelemetry-instrumentation-fastapi==0.55b1 \
opentelemetry-instrumentation-httpx==0.55b1 \
protobuf==5.28.3
# Copy the application code if needed
COPY ./amqp ./amqp
COPY ./otdemo ./otdemo
COPY ./demo.py .
# Create directory for uploaded files
RUN mkdir -p /tmp/clevermicro_documents
EXPOSE 8000 8001
CMD ["uvicorn", "demo:app", "--host", "0.0.0.0", "--port", "8000"]
+33
View File
@@ -116,3 +116,36 @@ AMQ_VERBOSE_LOGGING - if set to 1, the adapter will indicate in the log next to
As an open-source project we run entierly off donations. Buy one of our hardworking developers a beer by [donating at OpenCollective](https://opencollective.com/cleverthis). All donations go to our bounty fund and allow us to place bounties on important bugs and enhancements. You may also use Beerpay linked via the above badges.
## Running with Docker Compose
This project includes a Docker Compose configuration for easy local development and testing.
### Prerequisites
- Docker and Docker Compose installed on your system
### Starting the services
```bash
docker-compose up -d
```
This will start:
- The demo application (accessible at http://localhost:8000)
- Jaeger for distributed tracing (UI at http://localhost:16686)
- Prometheus for metrics collection (UI at http://localhost:9090)
- Alertmanager for alert handling (UI at http://localhost:9093)
- RabbitMQ for messaging (Management UI at http://localhost:15672)
### Stopping the services
```bash
docker-compose down
```
To remove volumes as well:
```bash
docker-compose down -v
```
+73 -164
View File
@@ -4,15 +4,17 @@ import time
from asyncio import AbstractEventLoop
from threading import Thread
import psutil
from aio_pika import Message
from aio_pika.abc import AbstractRobustChannel
from amqp.adapter.logging_utils import logging_info, logging_warning
from amqp.adapter.logging_utils import logging_info
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import ScalingRequestAlert
from amqp.router.utils import await_future, await_result
CPU_IDLE, CPU_ACTIVE, CPU_OVERLOAD = 0, 1, 2
IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 0.15, 0.90
RESOURCE_UNSET, RESOURCE_IDLE, RESOURCE_ACTIVE, RESOURCE_OVERLOAD = -1, 0, 1, 2
class ScaleRequestV1:
@@ -20,16 +22,12 @@ class ScaleRequestV1:
self,
serviceId: str,
taskId: str,
maxAvailability: int,
currentAvailability: int,
requestType: ScalingRequestAlert,
):
self.version = 1 # Version of the request, currently 1
self.serviceId = serviceId
self.taskId = taskId
self.maxAvailability = maxAvailability
self.currentAvailability = currentAvailability
self.requestType = requestType
self.thread = None
def to_json(self) -> str:
@@ -39,17 +37,22 @@ class ScaleRequestV1:
"version": self.version,
"serviceId": self.serviceId,
"taskId": self.taskId,
"maxAvailability": self.maxAvailability,
"currentAvailability": self.currentAvailability,
"requestType": self.requestType.value, # Using .value for Enum serialization
},
indent=2,
)
class RunningAverage:
def __init__(self, num_items):
self.buffer = [0] * num_items # Fixed-size array
"""
A class to maintain a running average of the last N values.
The intended use is to track CPU usage over time window, that is directly related to sample rate and window size.
This is the reason why the constructor takes the number of items to store, calculated as window size / sample rate
Initial_value is for unit test to emulate OVERLOAD or IDLE condition.
"""
def __init__(self, num_items, initial_value=0):
self.buffer = [initial_value] * num_items # Fixed-size array
self.pointer = 0 # Current position in array
self.num_items = num_items # Maximum number of items to store
self.is_filled = False # Track if buffer is fully filled
@@ -84,11 +87,12 @@ class RunningAverage:
class BackpressureHandler:
# Track the number of messages currently being processed
current_parallel_executions = 0
# helps detect IDLE condition
last_data_message_time = 0
# helps prevent flooding the system with backpressure events
last_backpressure_event_time = 0
last_backpressure_event = ScalingRequestAlert.UPDATE
# report when remain availability goes down
last_usage_report_time = 0
# report when remain availability goes up
last_recover_report_time = 0
# _callback_list is used in unit tests to record the invoked callbacks
_callback_list = None
def __init__(
self,
@@ -104,178 +108,79 @@ 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 increase_parallel_executions(self):
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
async def increase_parallel_executions(self):
"""Increase the number of parallel executions"""
logging_info(
"Backpressure: Increase parallel executions, current=%s",
self.current_parallel_executions,
)
self.current_parallel_executions += 1
def decrease_parallel_executions(self):
"""Decrease the number of parallel executions"""
logging_info(
"Backpressure: Decrease parallel executions, current=%s",
"Backpressure: Increased parallel executions, new value: %s",
self.current_parallel_executions,
)
# increase parallel -> reduce in remaining availability -> usage
# report usage at least every 1s
# TODO: add config for the interval
if self.last_usage_report_time + 1 < time.time():
await self.send_backpressure_report()
self.last_usage_report_time = time.time()
async def decrease_parallel_executions(self):
"""Decrease the number of parallel executions"""
if self.current_parallel_executions > 0:
self.current_parallel_executions -= 1
def update_parallel_executions(self, count: int):
"""Update the number of parallel executions"""
self.current_parallel_executions = count
def update_last_data_message_time(self):
logging_info("Backpressure: Update last data message time")
"""Update the last data message time"""
self.last_data_message_time = time.time()
logging_info(
"Backpressure: Decreased parallel executions, new value=%s",
self.current_parallel_executions,
)
# decrease parallel -> add to remaining availability -> recover
# report recover at every 3s
# TODO: add config for the interval
if self.last_recover_report_time + 3 < time.time():
await self.send_backpressure_report()
self.last_recover_report_time = time.time()
def start_backpressure_monitor(self) -> Thread:
# Start the Backpressure monitor loop
self.thread = Thread(target=self.backpressure_monitor_loop)
self.thread = Thread(target=self.regular_report_loop_wrapper)
self.thread.daemon = True # This makes it a daemon thread
self.thread.start()
return self.thread
async def check_overload_condition(self):
"""Check if the current parallel executions exceed the limit"""
self.update_last_data_message_time()
logging_info(
"Backpressure: Check overload condition, current=%s, max=%s",
self.current_parallel_executions,
self.parallel_workers,
)
if self.current_parallel_executions >= self.parallel_workers - 1:
# Check if the last backpressure event was not an overload
if self.last_backpressure_event != ScalingRequestAlert.OVERLOAD:
# Trigger the overload event
await self.handle_backpressure_overload_event()
self.last_backpressure_event_time = time.time()
self.last_backpressure_event = ScalingRequestAlert.OVERLOAD
async def _regular_report_loop(self):
"""Regularly report the availability in a fixed interval"""
while self._do_loop():
await self.send_backpressure_report()
# report recover at every 5s
# TODO: add config for the interval
await asyncio.sleep(5)
def backpressure_monitor_loop(self):
_time_window_millis = self.config.backpressure.time_window
_idle_duration_millis = self.config.backpressure.idle_duration
_overload_duration_millis = self.config.backpressure.idle_duration
def regular_report_loop_wrapper(self):
_loop = asyncio.new_event_loop()
_monitor_interval = 0.2 # Monitor every 200ms
CPU_IDLE, CPU_ACTIVE, CPU_OVERLOAD = 0, 1, 2
IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 15, 90
self.average_cpu_usage = RunningAverage(
int(max(_overload_duration_millis, _idle_duration_millis) // _monitor_interval)
)
_loop.run_until_complete(self._regular_report_loop())
async def _backpressure_monitor():
"""Monitor the backpressure conditions and trigger events accordingly"""
_cpu_usage_state = CPU_ACTIVE
_cpu_usage_changed = time.time()
while True:
_current_cpu_usage = psutil.cpu_percent(interval=None)
self.average_cpu_usage.add(_current_cpu_usage)
_average_cpu_usage = self.average_cpu_usage.get_average()
_now = time.time()
# logging_info(
# "Backpressure: current CPU[pct]=%s, avg=%s, state=%s, last dataMsg=%s, last event=%s, last eventTime=%s",
# _current_cpu_usage,
# _average_cpu_usage,
# _cpu_usage_state,
# self.last_data_message_time,
# self.last_backpressure_event,
# self.last_backpressure_event_time,
# )
if _average_cpu_usage < IDLE_THRESHOLD:
if _cpu_usage_state != CPU_IDLE:
_cpu_usage_changed = _now
_cpu_usage_state = CPU_IDLE
elif _average_cpu_usage > OVERLOAD_THRESHOLD:
if _cpu_usage_state != CPU_OVERLOAD:
_cpu_usage_changed = _now
_cpu_usage_state = CPU_OVERLOAD
else:
if _cpu_usage_state != CPU_ACTIVE:
_cpu_usage_changed = _now
_cpu_usage_state = CPU_ACTIVE
# Check if the current time exceeds the last backpressure event time
if (
_cpu_usage_state == CPU_OVERLOAD
and _now - _cpu_usage_changed > _overload_duration_millis
and (
_now - self.last_backpressure_event_time > _time_window_millis
or self.last_backpressure_event != ScalingRequestAlert.OVERLOAD
)
):
# Trigger the overload event
await self.handle_backpressure_overload_event()
self.last_backpressure_event_time = _now
self.last_backpressure_event = ScalingRequestAlert.OVERLOAD
elif (
_cpu_usage_state == CPU_IDLE
and _now - _cpu_usage_changed > _idle_duration_millis
and _now - self.last_data_message_time > _idle_duration_millis
and (
_now - self.last_backpressure_event_time > _time_window_millis
or self.last_backpressure_event != ScalingRequestAlert.IDLE
)
):
# Trigger the idle event
await self._handle_backpressure_idle_event()
self.last_backpressure_event = ScalingRequestAlert.IDLE
self.last_backpressure_event_time = _now
else:
# Trigger update event
if _now - self.last_backpressure_event_time > _time_window_millis:
# Trigger the update event
_scaling_request: ScaleRequestV1 = ScaleRequestV1(
self.swarm_service_id,
self.swarm_task_id,
100,
_average_cpu_usage,
ScalingRequestAlert.UPDATE,
)
# Address the message to any adapter capable of supporting BACKPRESSURE request
await self.publish_backpressure_request(_scaling_request)
self.last_backpressure_event_time = _now
self.last_backpressure_event = ScalingRequestAlert.UPDATE
await asyncio.sleep(_monitor_interval)
_loop.run_until_complete(_backpressure_monitor())
async def handle_backpressure_overload_event(self):
logging_warning("Backpressure: Capacity close to depleted!")
# Send an Overload event to the Management service
async def send_backpressure_report(self):
scaling_request: ScaleRequestV1 = ScaleRequestV1(
self.swarm_service_id,
self.swarm_task_id,
self.parallel_workers,
self.parallel_workers - self.current_parallel_executions,
ScalingRequestAlert.OVERLOAD,
serviceId=self.swarm_service_id,
taskId=self.swarm_task_id,
currentAvailability=self.parallel_workers - self.current_parallel_executions,
)
# Address the message to any management service instance capable of supporting BACKPRESSURE request
await self.publish_backpressure_request(scaling_request)
# update / reset time-window so that the OVERLOAD is not sent too often
self.last_data_message_time = time.time()
async def _handle_backpressure_idle_event(self):
logging_warning("Backpressure: Service is idle.")
# Send an Idle event to the Management service
scaling_request: ScaleRequestV1 = ScaleRequestV1(
self.swarm_service_id,
self.swarm_task_id,
self.parallel_workers,
self.parallel_workers - self.current_parallel_executions,
ScalingRequestAlert.IDLE,
)
# Address the message to any adapter capable of supporting BACKPRESSURE request
await self.publish_backpressure_request(scaling_request)
# update / reset time-window so that the IDLE is not sent too often
self.last_data_message_time = time.time()
async def publish_backpressure_request(self, scaling_request: ScaleRequestV1):
# Publish the backpressure request to the management service
logging_info(
f"Publishing backpressure for {scaling_request.serviceId} with request type {scaling_request.requestType}"
f"Publishing backpressure for {scaling_request.serviceId} with availability {scaling_request.currentAvailability}"
)
if not self.exchange:
@@ -283,6 +188,8 @@ class BackpressureHandler:
_exchange = await channel.get_exchange(name="cleverthis.clevermicro.management")
return _exchange
if BackpressureHandler._callback_list is not None:
BackpressureHandler._callback_list["_wrap_rabbit_mq_api_init"] = 1
self.exchange = await await_result(
asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api_init(self.channel), self.loop)
)
@@ -312,4 +219,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))
+166 -61
View File
@@ -4,6 +4,7 @@ Methods here will invoke the actual CleverThis Service code, which are passed as
"""
import asyncio
import enum
import inspect
import json
import pathlib
@@ -19,6 +20,7 @@ from aio_pika.abc import (
)
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
@@ -29,12 +31,45 @@ 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
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.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQErrorMessage, AMQMessage, 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]:
"""
Create objects for reply.
Return (channel for reply queue, exchange for reply, reply queue name)
"""
# Wrap calls to RabbitMQ API to make it Awaitable, so it can be run thread-safe with
# correct event loop. This is needed because the RabbitMQ API is not thread-safe and needs
# to be called in the context of event loop the first RabbitMQ connection was created with.
_channel = await connection.channel()
_exchange_name: str = AMQ_REPLY_TO_EXCHANGE
_exchange = await _channel.get_exchange(name=_exchange_name, ensure=True)
_queue = await _channel.declare_queue(exclusive=False) # Declare a unique, non-exclusive queue
_queue_name = _queue.name
await _queue.bind(
exchange=_exchange_name,
routing_key=_queue_name, # Use queue name as routing key
)
logging_debug(f"Publishing to exchange: {_exchange_name}, q: {_queue_name}")
return _channel, _exchange, _queue_name
async def _convert_file_response(
obj: FileResponse,
@@ -47,33 +82,12 @@ async def _convert_file_response(
The file content itself is transmitted through the named stream (which is implemented as
temporary dedicated queue.
"""
async def _wrap_amq_api_calls(
connection: AbstractRobustConnection,
) -> tuple[AbstractRobustChannel, AbstractExchange, str]:
# Wrap calls to RabbitMQ API to make it Awaitable, so it can be run thread-safe with
# correct event loop. This is needed because the RabbitMQ API is not thread-safe and needs
# to be called in the context of event loop the first RabbitMQ connection was created with.
_channel = await connection.channel()
_exchange_name: str = AMQ_REPLY_TO_EXCHANGE
_exchange = await _channel.get_exchange(name=_exchange_name, ensure=True)
_queue = await _channel.declare_queue(
exclusive=False
) # Declare a unique, non-exclusive queue
_queue_name = _queue.name
await _queue.bind(
exchange=_exchange_name,
routing_key=_queue_name, # Use queue name as routing key
)
logging_debug(f"Publishing to exchange: {_exchange_name}, q: {_queue_name}")
return _channel, _exchange, _queue_name
# All calls to RabbitMQ API have to be called threadsafe, because those appear to be attached
# to parent event loop, and need to execute at that event loop, not the current one
# allow coroutine to execute. calling _future.result() will block the entire thread
# and the coroutine will never execute
_channel, _exchange, _routing_key = await await_result(
asyncio.run_coroutine_threadsafe(_wrap_amq_api_calls(connection), loop)
asyncio.run_coroutine_threadsafe(_create_reply_exchange_and_queue(connection), loop)
)
# publish_file also needs RabbitMQ API, ensure the loop is passed as argument
(_stream_name, _file_size) = await file_handler.publish_file(
@@ -97,6 +111,8 @@ 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("]")
@@ -117,12 +133,17 @@ class CleverThisServiceAdapter:
amq_configuration: AMQConfiguration,
routes: List[APIRoute],
session: ClientSession = None,
authenticated_user_arg_name: str = "authenticated_user",
extra_init_data=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.
"""
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
@@ -132,6 +153,8 @@ class CleverThisServiceAdapter:
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:
"""
@@ -468,17 +491,21 @@ class CleverThisServiceAdapter:
_routes = self.routes.get(method, [])
_path, _args = AMQDataParser.parse_get_url(message.path())
for _route in _routes:
# TODO: here the regex will match with path prefix,
# but it might not be a stable feature that we could rely on.
# Need to figure out something stable for the prefix path.
_match = _route.path_regex.match(_path)
if _match:
_args.update(_match.groupdict())
_args["authenticated_user"] = auth_user
_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 CleverThisServiceAdapter.call_cleverthis_api(
return await self.call_cleverthis_api(
message,
_args,
api_method=_route.endpoint,
success_code=_route.status_code,
loop=self.loop,
)
return DataResponseFactory.of(
@@ -510,40 +537,45 @@ class CleverThisServiceAdapter:
_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]
_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"
(
_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["size"],
filename=_file_data["filename"],
headers=Headers({"Content-Type": _file_data["mediaType"]}),
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 /login input
if arg == "authenticated_user":
_verified_args["authenticated_user"] = auth_user
# 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 CleverThisServiceAdapter.call_cleverthis_api(
return await self.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 DataResponseFactory.of(
@@ -554,13 +586,98 @@ class CleverThisServiceAdapter:
message.trace_info(),
)
@staticmethod
def _verify_and_instantiate_arguments(
self,
args: Any,
api_method: Callable[[Any, Any], Coroutine],
) -> dict:
"""
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.
"""
_verified_args = {}
try:
_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)
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
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(
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,
loop: AbstractEventLoop | None = None,
) -> DataResponse:
"""
Handles POST and PUT requests, which typically contain possibly complex payload data,
@@ -575,31 +692,16 @@ class CleverThisServiceAdapter:
:return: DataResponse class
"""
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:
# convert to another type if arg is not a str
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] = serializer.str_to_bool(args[arg.name])
# TODO: more?
else:
_verified_args[arg.name] = args[arg.name]
else:
logging_info(f"{api_method.__name__}: Missing required argument: {arg.name}")
_verified_args = self._verify_and_instantiate_arguments(
args=args, api_method=api_method
)
logging_info(
"INVOKE ENDPOINT: {%s}, ARGS: {%s}",
"INVOKE ENDPOINT: [%s], ARGS: {%s}",
api_method.__name__,
", ".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 = (
@@ -607,7 +709,7 @@ class CleverThisServiceAdapter:
_result,
message.connection,
file_handler=StreamingFileHandler(),
loop=loop,
loop=self.loop,
),
"application/octet-stream",
)
@@ -615,7 +717,10 @@ class CleverThisServiceAdapter:
_json_result = _result.json()
_content_type = "application/json"
else:
_json_result = _result
if isinstance(_result, str):
_json_result = _result
else:
_json_result = python_type_to_json(_result)
_content_type = "application/json"
_binary_result = (
+103
View File
@@ -0,0 +1,103 @@
import os
import logging
import random
import socket
import time
from typing import Optional
import consul
from amqp.config.amq_configuration import AMQConfiguration
logger = logging.getLogger(__name__)
# Default configuration values for Adapter
DEFAULT_CONSUL_HOST = 'localhost'
DEFAULT_CONSUL_PORT = 8500
DEFAULT_CONSUL_COUNTER_KEY = 'adapter/ids/counter'
DEFAULT_MAX_RETRIES = 5
DEFAULT_RETRY_DELAY_SECONDS = 0.1
def get_config_values():
"""
Get configuration values from AMQConfiguration or environment variables.
Returns:
tuple: (consul_host, consul_port, consul_counter_key, max_retries, retry_delay_seconds)
"""
try:
config = AMQConfiguration("application.properties")
consul_host = os.environ.get('CONSUL_HOST', DEFAULT_CONSUL_HOST)
consul_port = int(os.environ.get('CONSUL_PORT', DEFAULT_CONSUL_PORT))
consul_counter_key = os.environ.get('CONSUL_COUNTER_KEY', DEFAULT_CONSUL_COUNTER_KEY)
max_retries = int(os.environ.get('CONSUL_MAX_RETRIES', DEFAULT_MAX_RETRIES))
retry_delay_seconds = float(os.environ.get('CONSUL_RETRY_DELAY_SECONDS', DEFAULT_RETRY_DELAY_SECONDS))
return consul_host, consul_port, consul_counter_key, max_retries, retry_delay_seconds
except Exception as e:
logger.warning(f"Failed to load configuration: {str(e)}. Using default values.")
return (
os.environ.get('CONSUL_HOST', DEFAULT_CONSUL_HOST),
int(os.environ.get('CONSUL_PORT', DEFAULT_CONSUL_PORT)),
os.environ.get('CONSUL_COUNTER_KEY', DEFAULT_CONSUL_COUNTER_KEY),
int(os.environ.get('CONSUL_MAX_RETRIES', DEFAULT_MAX_RETRIES)),
float(os.environ.get('CONSUL_RETRY_DELAY_SECONDS', DEFAULT_RETRY_DELAY_SECONDS))
)
def generate_fallback_id() -> int:
"""
Generate a fallback ID when Consul is not available.
Uses a combination of timestamp, hostname hash, and random number.
Returns:
int: A reasonably unique integer ID
"""
timestamp = int(time.time() * 1000)
hostname = socket.gethostname()
hostname_hash = hash(hostname) % 10000
random_part = random.randint(0, 9999)
unique_id = (timestamp * 100000000) + (hostname_hash * 10000) + random_part
logger.warning(f"Using fallback ID generation method: {unique_id}")
return unique_id
def get_unique_instance_id() -> int:
"""
Get a globally unique ID from Consul.
Uses Consul's atomic Compare-And-Set operations to safely increment a counter.
Falls back to a local generation method if Consul is unavailable.
Returns:
int: A globally unique integer ID
"""
consul_host, consul_port, consul_counter_key, max_retries, retry_delay_seconds = get_config_values()
try:
c = consul.Consul(host=consul_host, port=consul_port)
index, data = c.kv.get(consul_counter_key)
if data is None:
logger.info(f"Initializing Consul counter at {consul_counter_key}")
if c.kv.put(consul_counter_key, "1"):
return 1
else:
logger.error("Failed to initialize Consul counter")
return generate_fallback_id()
current_value = int(data['Value'].decode('utf-8'))
new_value = current_value + 1
for attempt in range(max_retries):
success = c.kv.put(
consul_counter_key,
str(new_value),
cas=data['ModifyIndex']
)
if success:
logger.debug(f"Successfully obtained unique ID: {new_value}")
return new_value
logger.debug(f"CAS update failed on attempt {attempt+1}, retrying...")
time.sleep(retry_delay_seconds)
index, data = c.kv.get(consul_counter_key)
if data is None:
logger.error("Counter disappeared during update")
return generate_fallback_id()
current_value = int(data['Value'].decode('utf-8'))
new_value = current_value + 1
logger.error(f"Failed to update counter after {max_retries} attempts")
return generate_fallback_id()
except Exception as e:
logger.error(f"Error accessing Consul: {str(e)}")
return generate_fallback_id()
+28 -1
View File
@@ -7,10 +7,11 @@ import json
import re
from datetime import datetime, timezone
from io import BytesIO
from typing import Dict, List
from typing import Any, Dict, List
import amqp.adapter.file_handler
from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
from amqp.adapter.pydantic_serializer import python_type_to_json
from amqp.adapter.serializer import (
map_as_string,
map_with_list_as_string,
@@ -112,6 +113,32 @@ class DataMessageFactory:
payload,
)
def create_rpc_message(
self,
fq_method_name: str,
swarm_id: str,
args: Dict[str, Any] = None,
) -> DataMessage:
"""
Create a RPC message with the given parameters.
:param fq_method_name: fully qualified method name (e.g. 'com.example.ServiceClass.method_name')
:param args: arguments whose serialized form will form the request body
:return: DataMessage object
"""
package, class_name, method = fq_method_name.rsplit(".", 2)
headers = {"Content-Type": ["application/octet-stream"], "Return-Address": [swarm_id]}
trace_info = {}
return DataMessage(
DataMessageMagicByte.RPC_REQUEST.value,
self.generate_snowflake_id(),
package,
class_name,
method,
headers,
trace_info,
python_type_to_json(args) if args else b"",
)
@staticmethod
def amq_message_routing_key(message: DataMessage) -> str:
"""
+6 -6
View File
@@ -40,7 +40,7 @@ class MultipartFormDataParser:
try:
self.form_data = json.loads(message.body())
self.is_json = True
if self.form_data["files"] and self.form_data["form"]:
if self.form_data["files"] is not None and self.form_data["form"] is not None:
self.form_data = AMQDataParser.process_form_data(self.form_data)
except Exception as jsonerror:
logging_error(f"Parsing JSON: {jsonerror}")
@@ -51,13 +51,14 @@ class MultipartFormDataParser:
)
self.is_multipart = True
except Exception as e:
logging_error(f"Error parsing multipart/form-data: {e}. Trying parsing as JSON.")
__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"] is not None and self.form_data["form"] is not None:
self.form_data = AMQDataParser.process_form_data(self.form_data)
except Exception as jsonerror:
logging_error(__error)
logging_error(f"Parsing JSON: {jsonerror}")
self.is_valid = False
@@ -109,7 +110,6 @@ class AMQDataParser:
for key, value in data["files"].items():
if value: # Only override if files data is non-empty
result[key] = value[0]
return result
@staticmethod
@@ -149,7 +149,7 @@ class AMQDataParser:
# Parse JSON
if content_type == "application/json":
try:
return json.loads(message.body())
return json.loads(message.body() or "{}")
except json.JSONDecodeError as e:
logging_error(f"Invalid JSON: {e}")
raise ValueError(f"Invalid JSON: {e}")
@@ -162,7 +162,7 @@ class AMQDataParser:
}
if len(_form_data) == 0 and len(message.body()) > 0:
_form_data = json.loads(message.body())
if _form_data["files"] and _form_data["form"]:
if _form_data["files"] is not None and _form_data["form"] is not None:
_form_data = AMQDataParser.process_form_data(_form_data)
return _form_data
except Exception as e:
+74 -50
View File
@@ -153,6 +153,44 @@ class StreamingFileHandler(FileHandler):
def __init__(self):
super().__init__()
# This method MUST NOT be static.
# For testing purpose, we will mock this method and inject a mock connection.
async def _get_connection(self, rabbitmq_url, loop):
return await connect_robust(rabbitmq_url, loop=loop)
async def _download_from_queue_to_file(
self, rabbitmq_url, queue_name, filepath, loop
) -> tuple[int, int] | None:
"""
Internal method for downloading a file from a queue to disk.
Return (file size, eof status).
"""
# Here the normal mock doesn't work since this method
# is async and will be run in another loop,
# out of the scope of method mock.
connection = await self._get_connection(rabbitmq_url, loop=loop)
_file_size = 0
_eof = IN_PROGRESS
try:
channel: AbstractRobustChannel = await connection.channel()
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
logging_debug(f"Awaiting all file chunks being downloaded for {filepath}")
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
with open(filepath, "ab") as f:
f.write(message.body)
_file_size = _file_size + message.body_size
_eof = int(message.headers.get("eof", END_OF_FILE))
if _eof != IN_PROGRESS:
break
except Exception as e:
logging_error(f"File download failed: {e}")
finally:
await connection.close()
return _file_size, _eof
async def download_file(
self, loop, rabbitmq_url, file_info, message_data, output_dir
) -> (int, int):
@@ -184,39 +222,46 @@ class StreamingFileHandler(FileHandler):
f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name}:{queue_name} to {filepath}"
)
async def wrap_rabbit_mq_calls() -> tuple[int, int] | None:
connection = await connect_robust(rabbitmq_url, loop=loop)
try:
_file_size = 0
_eof = IN_PROGRESS
channel: AbstractRobustChannel = await connection.channel()
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
logging_debug(f"Awaiting all file chunks being downloaded for {filename}")
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
with open(filepath, "ab") as f:
f.write(message.body)
_file_size = _file_size + message.body_size
_eof = int(message.headers.get("eof", END_OF_FILE))
if _eof != IN_PROGRESS:
break
except Exception as e:
logging_error(f"File download failed: {e}")
finally:
await connection.close()
return _file_size, _eof
if queue_name:
_local_file_size, _local_eof = await await_result(
asyncio.run_coroutine_threadsafe(wrap_rabbit_mq_calls(), loop=loop)
asyncio.run_coroutine_threadsafe(
self._download_from_queue_to_file(rabbitmq_url, queue_name, filepath, loop),
loop=loop,
)
)
return _local_file_size, _local_eof
return 0, CANCELLED
async def download_buffer(self, loop, rabbitmq_url, queue_name) -> (bytearray, int):
async def _download_from_queue_to_buffer(
self, rabbitmq_url, queue_name, loop
) -> tuple[bytearray, int]:
"""
Internal method for downloading a file from a queue to RAM.
Return (bytearray, eof status).
"""
_connection = await self._get_connection(rabbitmq_url, loop=loop)
_eof = IN_PROGRESS
_res = bytearray()
try:
channel: AbstractRobustChannel = await _connection.channel()
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
logging_debug(f"Awaiting all buffer chunks being downloaded for {queue_name}")
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
_res.extend(message.body)
_eof = message.headers.get("eof", END_OF_FILE)
if _eof != IN_PROGRESS:
break
logging_info(f"Finished downloading buffer {queue_name}")
except Exception as e:
logging_error(f"Downloading buffer: {e}")
await _connection.close()
return _res, _eof
async def download_buffer(self, loop, rabbitmq_url, queue_name: str | None) -> (bytearray, int):
"""
Asynchronously downloads a bytearay from RabbitMQ based on the provided file information.
It assumes entire array is in the first (and only) chunk
@@ -227,31 +272,10 @@ class StreamingFileHandler(FileHandler):
"""
logging_debug(f"Downloading buffer from stream: {queue_name}")
if queue_name:
async def wrap_rabbit_mq_calls() -> tuple[bytearray, int]:
_connection = await connect_robust(rabbitmq_url, loop=loop)
_eof = IN_PROGRESS
_res = bytearray()
try:
channel: AbstractRobustChannel = await _connection.channel()
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
logging_debug(f"Awaiting all buffer chunks being downloaded for {queue_name}")
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
_res.extend(message.body)
_eof = message.headers.get("eof", END_OF_FILE)
if _eof != IN_PROGRESS:
break
logging_info(f"Finished downloading buffer {queue_name}")
except Exception as e:
logging_error(f"Downloading buffer: {e}")
await _connection.close()
return _res, _eof
_local_res, _local_eof = await await_result(
asyncio.run_coroutine_threadsafe(wrap_rabbit_mq_calls(), loop=loop)
asyncio.run_coroutine_threadsafe(
self._download_from_queue_to_buffer(rabbitmq_url, queue_name, loop), loop=loop
)
)
return _local_res, _local_eof
return 0, CANCELLED
+1 -1
View File
@@ -120,7 +120,7 @@ def logging_debug(msg: str, *args) -> None:
_thread_id = threading.get_ident()
module, line_number, function_name = get_context_info()
if module and line_number and function_name:
logging.info(
logging.debug(
f"{_thread_id}:{__version__} [DBG] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
*args,
)
+90 -81
View File
@@ -2,10 +2,12 @@
This module provides serialization and deserialization functions for Pydantic models that extend BaseModel
"""
import importlib
import json
import uuid
from datetime import datetime
from typing import Any, Dict
from 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
from pydantic import BaseModel
@@ -26,8 +28,12 @@ def serialize_object(obj: BaseModel) -> str:
A string representation of the object.
"""
class_name = obj.__class__.__name__
module_name = obj.__class__.__module__
obj_dict = {}
for key, value in obj.model_dump().items():
for key, _ in obj.model_dump().items():
# here we use the original value
# the one returned by model_dump will convert everything into dict recursively
value = getattr(obj, key)
if isinstance(value, uuid.UUID):
obj_dict[key] = str(value)
elif isinstance(value, datetime):
@@ -47,8 +53,50 @@ def serialize_object(obj: BaseModel) -> str:
else:
obj_dict[key] = value
dict_string = ",".join(f"{k}:{v}" for k, v in obj_dict.items())
return f"{class_name}:{{{dict_string}}}"
dict_string = json.dumps(obj_dict)
return f"{module_name}.{class_name}:{dict_string}"
def 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
def deserialize_object(data: str) -> BaseModel:
@@ -63,93 +111,54 @@ def deserialize_object(data: str) -> BaseModel:
Returns:
An R2RSerializable object.
"""
class_name, dict_str = data.split(":", 1)
dict_str = dict_str.strip("{}")
pairs = dict_str.split(",") if dict_str else []
class_path, dict_str = data.split(":", 1)
module_path, class_name = class_path.rsplit(".", 1)
obj_dict: Dict[str, Any] = {}
for pair in pairs:
if not pair or ":" not in pair:
continue
k, v = pair.split(":", 1)
v = v.strip() # important
obj_dict[k] = _convert_value(v)
obj_dict: Dict[str, Any] = json.loads(dict_str)
logging_debug(f"Deserialized object: {class_name}")
logging_debug(f"Deserialized object: {class_path}")
if "[" in class_name:
class_name = class_name.split("[")[0]
logging_debug(f"trying object: {class_name}")
cls = globals().get(class_name)
module = importlib.import_module(module_path)
cls = getattr(module, class_name)
if cls is None:
raise ValueError(f"Unknown class: {class_name}")
# process nested model, dict and list
for key, value in obj_dict.items():
key_type: type[Any] = cls.model_fields[key].annotation
# handle union with none, this will give the true type
# for example, return `A` if the type is `A | None`
if get_origin(key_type) is UnionType or get_origin(key_type) is Union:
union_types = get_args(key_type)
if len(union_types) > 2:
# if there are multiple unions, like `A | B | None`
# then we have no way to figure out which one to use
raise ValueError(f"Unsupported Union type: {key_type}")
# return the none-None one as the type for processing
if union_types[0] is type(None):
key_type = union_types[1]
else:
key_type = union_types[0]
# skip if the value is None
if value is None:
continue
if key_type == uuid.UUID:
obj_dict[key] = uuid.UUID(value)
elif key_type == datetime:
obj_dict[key] = datetime.fromisoformat(value)
elif key_type == dict:
obj_dict[key] = value
elif issubclass(key_type, BaseModel):
obj_dict[key] = deserialize_object(value)
elif key_type == list:
obj_dict[key] = list(map(try_deserialize_object, value))
return cls(**obj_dict)
def _convert_value(v: str) -> Any:
"""
Helper function to convert a string to its appropriate Python type.
"""
if v == "None":
return None
elif _is_uuid_string(v):
return uuid.UUID(v)
elif _is_datetime_string(v):
return datetime.fromisoformat(v)
elif _is_int_string(v):
return int(v)
elif (
v.startswith("{")
and v.endswith("}")
and ":" in v
and "{" not in v[1:-1]
and "}" not in v[1:-1]
): # nested object
# check if it is a json dict
try:
json.loads(v)
return json.loads(v)
except json.JSONDecodeError:
return deserialize_object(v) # recursive call
elif v.startswith("[") and v.endswith("]"): # list
v = v[1:-1] # remove brackets
items = v.split(",") if v else []
deserialized_list = []
for item in items:
item = item.strip()
deserialized_list.append(_convert_value(item)) # recursive call
return deserialized_list
else:
return v # Keep as string
def _is_uuid_string(s: str) -> bool:
"""Helper to check if a string is a valid UUID."""
try:
uuid.UUID(s)
return True
except ValueError:
return False
def _is_datetime_string(s: str) -> bool:
"""Helper to check if a string is a valid datetime in ISO format"""
try:
datetime.fromisoformat(s)
return True
except ValueError:
return False
def _is_int_string(s: str) -> bool:
"""Helper to check if a string is a valid int."""
try:
int(s)
return True
except ValueError:
return False
def parse_url_query(url):
"""
Parses a HTTP GET URL and returns (context_path, query_params_dict)
+5 -2
View File
@@ -18,7 +18,7 @@ def bytes_to_long(b: bytes) -> int:
def read_long(bis: BytesIO) -> int:
try:
return struct.unpack("q", bis.read(8))[0]
except IOError:
except struct.error:
return 0
@@ -51,7 +51,10 @@ def csv_as_list(input_str: str) -> List[str]:
def add_to_map(map_: Dict[str, List[str]], token: str) -> None:
eq_idx = token.index("=")
try:
eq_idx = token.index("=")
except ValueError:
eq_idx = -1 # set to -1 when not found
if eq_idx > 0:
key = token[:eq_idx]
map_[key] = csv_as_list(token[eq_idx + 1 :])
+205
View File
@@ -0,0 +1,205 @@
""" """
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
+5
View File
@@ -40,6 +40,11 @@ class AMQAdapter:
self.PREFIX + "require-authenticated-user",
fallback=False,
)
self.backpressure_enabled = config.getboolean(
"CleverMicro-AMQ",
self.PREFIX + "default-backpressure-enabled",
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_task_id = os.getenv("SWARM_TASK_ID", default="python-adapter")
+4 -2
View File
@@ -14,6 +14,8 @@ cm.amq-adapter.generator-id=1
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.require-authenticated-user=false
# Indicate if AMQ Adapter should use inbuilt backpressure monitor (True) or rely only on the service backpressure updates (False)
cm.amq-adapter.default-backpressure-enabled=false
# CleverMicro AMQ Adapter dispatcher, settings related to RabbitMQ connection and optional REST forwarding
cm.dispatch.use-dlq=true
@@ -27,7 +29,8 @@ cm.dispatch.download-dir=/tmp/downloaded_files
#cm.dispatch.service-port=8080
# Backpressure monitor settings
# Define maximum number of messages of being processed in parallel before triggering backpressure OVERLOAD alert
# Define maximum number of resources that service can use before triggering backpressure OVERLOAD alert
# 'resource' is a generic term that can mean CPU, memory, parallel requests, etc.
cm.backpressure.threshold=5
# Define maximum CPU usage (%) before triggering backpressure OVERLOAD alert
cm.backpressure.threshold-cpu-overload=90
@@ -39,4 +42,3 @@ cm.backpressure.time-window=10
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
-12
View File
@@ -1,12 +0,0 @@
import sys
import requests
url = "http://localhost:8080/amq-adapter-healthcheck"
headers = {"Authorization": "Bearer my_access_token"}
params = {"limit": 10, "offset": 20}
response = requests.get(url, headers=headers, params=params)
if response.status_code == 200:
sys.exit(0)
# Failed execution
sys.exit(1)
+3 -14
View File
@@ -91,6 +91,9 @@ class CleverMicroMessage:
def content_type(self) -> str:
return self.headers().get("Content-Type", self.DEFAULT_CONTENT_TYPE)[0]
def return_address(self) -> str:
return self.headers().get("Return-Address", "")[0]
class DataMessage(CleverMicroMessage):
"""
@@ -288,21 +291,11 @@ class ServiceMessage:
return self.message_type is not None
class ScalingRequestAlert(Enum):
OVERLOAD = "OVERLOAD"
IDLE = "IDLE"
UPDATE = "UPDATE"
SCALE_UP = "SCALE_UP"
SCALE_DOWN = "SCALE_DOWN"
@dataclass(frozen=True)
class ScalingRequest:
service_id: str
task_id: str
max_availability: int
current_availability: int
request_type: ScalingRequestAlert
version: int = 1
def to_json(self) -> str:
@@ -312,9 +305,7 @@ class ScalingRequest:
"version": self.version,
"serviceId": self.service_id,
"taskId": self.task_id,
"maxAvailability": self.max_availability,
"currentAvailability": self.current_availability,
"requestType": self.request_type.value, # Using .value for Enum serialization
},
indent=2,
)
@@ -326,9 +317,7 @@ class ScalingRequest:
return cls(
service_id=data["serviceId"],
task_id=data["taskId"],
max_availability=data["maxAvailability"],
current_availability=data["currentAvailability"],
request_type=ScalingRequestAlert(data["requestType"]), # Convert string to Enum
)
+2 -3
View File
@@ -17,7 +17,6 @@ from amqp.router.router_base import (
AMQ_RPC_EXCHANGE,
DLQ_EXCHANGE,
)
from amqp.router.router_consumer import RouterConsumer
from amqp.router.utils import sanitize_as_rabbitmq_name
@@ -25,7 +24,7 @@ class RabbitMQClient:
PREFETCH_COUNT = 1
DLQ_ARGS = {"x-dead-letter-exchange": DLQ_EXCHANGE}
def __init__(self, configuration):
def __init__(self, configuration, router):
self.channel: aio_pika.RobustChannel | None = None
self.connection: aio_pika.RobustConnection | None = None
self.rpc_exchange: aio_pika.RobustExchange | None = None
@@ -37,7 +36,7 @@ class RabbitMQClient:
self.amq_configuration.dispatch.amq_host,
)
logging.info("*******************************************")
self.router = RouterConsumer(configuration)
self.router = router
# this.meterRegistry = new PrometheusMeterRegistry(PrometheusConfig.DEFAULT);
self.reply_to_exchange: AbstractRobustExchange | None = (
+8 -1
View File
@@ -1,5 +1,5 @@
import re
from typing import Set
from typing import Optional, Set
from amqp.adapter.logging_utils import logging_info
from amqp.model.model import AMQRoute
@@ -55,6 +55,13 @@ class RouteDatabase:
return route
return None
def find_route_by_service(self, service_name: str) -> Optional[AMQRoute]:
logging_info(f"findRoute service_name:{service_name} in:{self.wrap_routes()}")
for route in self.routes:
if route.component_name == service_name:
return route
return None
def add_route(self, route: AMQRoute) -> bool:
if route not in self.routes:
self.routes.add(route)
+4 -1
View File
@@ -1,4 +1,4 @@
from typing import Callable, List, Set
from typing import Callable, List, Optional, Set
from amqp.adapter.amq_route_factory import NULL_ROUTE, AMQRouteFactory
from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
@@ -49,6 +49,9 @@ class RouterBase:
def find_route_by_message(self, message: DataMessage) -> AMQRoute | None:
return self.route_database.find_route(message.domain(), message.path())
def find_route_by_service(self, service_name: str) -> Optional[AMQRoute]:
return self.route_database.find_route_by_service(service_name)
def add_routes(self, message: str) -> None:
try:
amq_route_list = self.unwrap_route_list(message)
+3 -4
View File
@@ -124,7 +124,7 @@ class DataMessageHandler:
:return: DataResponse
"""
# Increment the counter for parallel executions
self.backpressure_handler.increase_parallel_executions()
await self.backpressure_handler.increase_parallel_executions()
logging_info(
f"PARALLEL: Current Capacity [{self.backpressure_handler.current_parallel_executions}] of [{self.backpressure_handler.parallel_workers}]"
)
@@ -132,7 +132,6 @@ class DataMessageHandler:
self.reply_to = message.reply_to
try:
await self.backpressure_handler.check_overload_condition()
_a_message = AMQMessage(amq_message, self.rabbit_mq_client.connection)
_response: DataResponse = await self.service_adapter.on_message(_a_message)
logging_debug(f"Result in thread: {_response}")
@@ -154,8 +153,8 @@ class DataMessageHandler:
logging_error(f"Deleting file {filename}: {e}")
except Exception as e:
logging_error(f"Request handler: {e}")
self.backpressure_handler.decrease_parallel_executions()
# recover availability
await self.backpressure_handler.decrease_parallel_executions()
async def reconstitute_data_message(
self, message: AbstractIncomingMessage
+55 -9
View File
@@ -14,10 +14,11 @@ from amqp.adapter.data_message_factory import DataMessageFactory
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.model.model import AMQRoute, DataMessage
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient
from amqp.router.router_consumer import RouterConsumer
from amqp.service.amq_message_handler import DataMessageHandler
from amqp.service.tracing import initialize_trace
from amqp.service.tracing import TracingConfig, initialize_telemetry
logging.basicConfig(level=logging.DEBUG)
logging.info(f"CWD = {os.getcwd()}")
@@ -45,15 +46,30 @@ class AMQService:
self.data_message_factory = DataMessageFactory.get_instance(
self.amq_configuration.amq_adapter.generator_id
)
self.rabbit_mq_client = RabbitMQClient(self.amq_configuration)
self.tracer = initialize_trace(name=amq_configuration.amq_adapter.service_name)
self.router = RouterConsumer(self.amq_configuration)
self.rabbit_mq_client = RabbitMQClient(self.amq_configuration, self.router)
self.tracer = initialize_telemetry(
tracing_config=TracingConfig(
service_name=amq_configuration.amq_adapter.service_name,
service_version=os.getenv("SERVICE_VERSION", "unknown"),
)
)
self.amq_data_message_handler = None
self.service_message_factory = ServiceMessageFactory(
self.amq_configuration.amq_adapter.service_name
)
self.backpressure_thread: Optional[Thread] = None
self.backpressure_handler: Optional[BackpressureHandler] = None
# =======================================================================
# =======================================================================
# ====================== A P I M e t h o d s =========================
# =======================================================================
# =======================================================================
def run(self):
"""
Start the AMQ service, initializing RabbitMQ connection and message handlers.
"""
# self.init(loop) will initialize RabbitMQ connection
asyncio.set_event_loop(self.loop)
self.loop.run_until_complete(self.init(self.loop, self.service_adapter))
@@ -63,9 +79,38 @@ class AMQService:
logging.info("***********************************************")
self.loop.run_forever()
def rpc(self, service_name: str, fq_method_name: str, **kwargs) -> Future:
"""
Send a Remote Procedure Call (RPC) message to the AMQ service.
:param service_name: The service name to receive the RPC message.
:param fq_method_name: The fully qualified method name to execute the RPC message.
:param kwargs: Additional keyword arguments to be passed to the RPC method.
:return: A Future that resolves when the RPC response is received.
"""
route: Optional[AMQRoute] = self.router.find_route_by_service(service_name)
if route is not None:
message: DataMessage = self.data_message_factory.create_rpc_message(
fq_method_name=fq_method_name,
swarm_id=self.amq_configuration.amq_adapter.swarm_task_id,
**kwargs,
)
result: Future = self.amq_data_message_handler.send_rpc(route, message)
self.set_outstanding_future(message, result)
return result
result = asyncio.Future()
result.set_exception(
NameError(f"No route exists. Service '{service_name}' not found in routing table.")
)
return result
# =======================================================================
# ====================== Internal Methods ============================
# =======================================================================
async def init(self, loop, service_adapter):
_reply_to_exchange: AbstractRobustExchange = await self.rabbit_mq_client.init(loop)
backpressure_handler = BackpressureHandler(
self.backpressure_handler = BackpressureHandler(
channel=self.rabbit_mq_client.channel,
loop=loop,
config=self.amq_configuration,
@@ -79,11 +124,12 @@ class AMQService:
self.rabbit_mq_client,
self.amq_configuration,
loop=loop,
backpressure_handler=backpressure_handler,
backpressure_handler=self.backpressure_handler,
)
await self.register_routes()
await self.rabbit_mq_client.request_routes()
self.backpressure_thread = backpressure_handler.start_backpressure_monitor()
logging_info("Starting backpressure monitor thread")
self.backpressure_thread = self.backpressure_handler.start_backpressure_monitor()
def cleanup(self):
logging_info("AMQService.cleanup() rmqc=%s", self.rabbit_mq_client)
@@ -92,7 +138,7 @@ class AMQService:
async def reinitialize_if_disconnected(self):
if not self.rabbit_mq_client.channel or self.rabbit_mq_client.channel.is_closed:
self.rabbit_mq_client = RabbitMQClient(self.amq_configuration)
self.rabbit_mq_client = RabbitMQClient(self.amq_configuration, self.router)
await self.rabbit_mq_client.init(loop=self.loop)
await self.register_routes()
await self.rabbit_mq_client.request_routes()
@@ -117,7 +163,7 @@ class AMQService:
)
return None
def send_message(self, message: DataMessage, future: Future):
def set_outstanding_future(self, message: DataMessage, future: Future):
self.amq_data_message_handler.outstanding[str(message.id())] = future
def remove_outstanding_future(self, message_id: str):
+233
View File
@@ -0,0 +1,233 @@
import consul_kv
import socket
import os
import time
import logging
import threading
# --- Configuration ---
CONSUL_HOST = os.environ.get('CONSUL_HOST', 'localhost')
CONSUL_PORT = int(os.environ.get('CONSUL_PORT', 8500))
CONSUL_COUNTER_KEY = 'service/ids/counter'
MAX_RETRIES = 5 # Max attempts for atomic update
RETRY_DELAY_SECONDS = 0.1 # Delay between retries for atomic update
# --- Logging Setup ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(threadName)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# --- Fallback ID Generation ---
def generate_fallback_id() -> int:
"""
Generates a fallback ID based on hostname and PID if Consul is unavailable.
This provides a unique-ish ID for local testing/development but is not guaranteed
to be globally unique across different machines or after restarts.
"""
hostname = socket.gethostname()
pid = os.getpid()
# A simple way to combine them into a number.
# Hash hostname to an integer and combine with PID.
# Note: This is a simple fallback. For production, consider a more robust
# local ID generation or a different strategy if Consul is critical.
fallback_id = (hash(hostname) % 1000000) * 100000 + pid
logger.warning(
f"Consul is unavailable. Generating fallback ID: {fallback_id} (based on hostname: {hostname}, PID: {pid})")
return fallback_id
# --- Unique Instance ID Generation Function ---
def get_unique_instance_id() -> int:
"""
Reads and atomically increments a counter in Consul KV to get a unique instance ID.
If Consul is unavailable or the atomic update fails after retries,
it falls back to generating an ID based on hostname and PID.
"""
try:
# Initialize Consul client
client = consul_kv.Connection(endpoint=f"{CONSUL_HOST}:{CONSUL_PORT}")
logger.info(f"Attempting to connect to Consul at {CONSUL_HOST}:{CONSUL_PORT}")
for attempt in range(MAX_RETRIES):
try:
# 1. Read the current value and its modifyIndex (for CAS)
# client.get returns (value, modifyIndex)
current_value_str, modify_index = client.get(CONSUL_COUNTER_KEY)
# If key doesn't exist, initialize it to 0
if current_value_str is None:
current_value = 0
modify_index = 0 # For a non-existent key, modify_index is 0 for initial CAS
logger.info(f"Consul key '{CONSUL_COUNTER_KEY}' not found. Initializing to 0.")
else:
current_value = int(current_value_str)
new_value = current_value + 1
# 2. Attempt atomic update using CAS (Check-And-Set)
# client.set returns True on success, False on failure (if modify_index changed)
success = client.put(CONSUL_COUNTER_KEY, str(new_value), cas=modify_index)
if success:
logger.info(f"Successfully obtained unique ID: {new_value} (Attempt {attempt + 1}/{MAX_RETRIES})")
return new_value
else:
logger.warning(
f"CAS failed for '{CONSUL_COUNTER_KEY}' (modify_index changed). Retrying... (Attempt {attempt + 1}/{MAX_RETRIES})")
time.sleep(RETRY_DELAY_SECONDS * (attempt + 1)) # Exponential backoff for retries
except ValueError:
# Handle case where the stored value is not an integer
logger.error(
f"Value stored at '{CONSUL_COUNTER_KEY}' is not a valid integer. Please check Consul KV. Attempting to overwrite.")
# To recover, you might force a set without CAS, but that risks data loss.
# For this demo, we'll just let it retry or fall back.
time.sleep(RETRY_DELAY_SECONDS * (attempt + 1))
except Exception as e:
# Catch specific ConsulKV errors (e.g., connection issues, invalid response)
logger.error(f"ConsulKV error during atomic update (Attempt {attempt + 1}/{MAX_RETRIES}): {e}")
time.sleep(RETRY_DELAY_SECONDS * (attempt + 1)) # Exponential backoff
logger.error(f"Failed to obtain unique ID from Consul after {MAX_RETRIES} attempts.")
return generate_fallback_id() # Fallback if all retries fail
except Exception as e:
# Catch broader connection errors or unexpected issues with Consul client initialization
logger.error(f"Could not connect to Consul or an unexpected error occurred: {e}")
return generate_fallback_id() # Fallback if Consul is completely unreachable
import os
import logging
import random
import socket
import time
from typing import Optional
import consul
from amqp.config.amq_configuration import AMQConfiguration
from amqp.config.amq_configuration import AMQConfiguration
# Configuration from environment variables with defaults
CONSUL_HOST = os.environ.get('CONSUL_HOST', 'localhost')
CONSUL_PORT = int(os.environ.get('CONSUL_PORT', 8500))
CONSUL_COUNTER_KEY = 'service/ids/counter'
MAX_RETRIES = 5 # Max attempts for atomic update
RETRY_DELAY_SECONDS = 0.1 # Delay between retries for atomic update
logger = logging.getLogger(__name__)
def generate_fallback_id() -> int:
"""
Generate a fallback ID when Consul is not available.
Uses a combination of timestamp, hostname hash, and random number.
Returns:
int: A reasonably unique integer ID
"""
# Get current timestamp in milliseconds
timestamp = int(time.time() * 1000)
# Get hostname and hash it to an integer
hostname = socket.gethostname()
hostname_hash = hash(hostname) % 10000 # Limit to 4 digits
# Generate a random number
random_part = random.randint(0, 9999) # 4 digits
# Combine all parts into a single integer
# Format: TTTTTTTTTTT-HHHH-RRRR (T=timestamp, H=hostname hash, R=random)
unique_id = (timestamp * 100000000) + (hostname_hash * 10000) + random_part
logger.warning(f"Using fallback ID generation method: {unique_id}")
return unique_id
def get_unique_instance_id() -> int:
"""
Get a globally unique ID from Consul.
Uses Consul's atomic Compare-And-Set operations to safely increment a counter.
Falls back to a local generation method if Consul is unavailable.
Returns:
int: A globally unique integer ID
"""
consul_host, consul_port, consul_counter_key, max_retries, retry_delay_seconds = get_config_values()
try:
# Connect to Consul
c = consul.Consul(host=consul_host, port=consul_port)
# Try to get the current counter value
index, data = c.kv.get(CONSUL_COUNTER_KEY)
# If the key doesn't exist yet, create it with initial value
if data is None:
logger.info(f"Initializing Consul counter at {CONSUL_COUNTER_KEY}")
if c.kv.put(CONSUL_COUNTER_KEY, "1"):
return 1
else:
logger.error("Failed to initialize Consul counter")
return generate_fallback_id()
# Get the current value
current_value = int(data['Value'].decode('utf-8'))
new_value = current_value + 1
# Try to atomically update the counter
for attempt in range(MAX_RETRIES):
# Use Compare-And-Set to ensure atomicity
success = c.kv.put(
CONSUL_COUNTER_KEY,
str(new_value),
cas=data['ModifyIndex']
)
if success:
logger.debug(f"Successfully obtained unique ID: {new_value}")
return new_value
# If CAS failed, someone else updated the value, retry
logger.debug(f"CAS update failed on attempt {attempt+1}, retrying...")
time.sleep(RETRY_DELAY_SECONDS)
# Get the latest value for the next attempt
index, data = c.kv.get(CONSUL_COUNTER_KEY)
if data is None:
logger.error("Counter disappeared during update")
return generate_fallback_id()
current_value = int(data['Value'].decode('utf-8'))
new_value = current_value + 1
# If we've exhausted all retries, use the fallback
logger.error(f"Failed to update counter after {MAX_RETRIES} attempts")
return generate_fallback_id()
except Exception as e:
logger.error(f"Error accessing Consul: {str(e)}")
return generate_fallback_id()
def get_config_values():
"""
Get configuration values from AMQConfiguration or environment variables.
Returns:
tuple: (consul_host, consul_port, consul_counter_key, max_retries, retry_delay_seconds)
"""
try:
config = AMQConfiguration("application.properties")
consul_host = os.environ.get('CONSUL_HOST', 'localhost')
consul_port = int(os.environ.get('CONSUL_PORT', 8500))
consul_counter_key = os.environ.get('CONSUL_COUNTER_KEY', 'service/ids/counter')
max_retries = int(os.environ.get('CONSUL_MAX_RETRIES', 5))
retry_delay_seconds = float(os.environ.get('CONSUL_RETRY_DELAY_SECONDS', 0.1))
return consul_host, consul_port, consul_counter_key, max_retries, retry_delay_seconds
except Exception as e:
logger.warning(f"Failed to load configuration: {str(e)}. Using default values.")
return (
os.environ.get('CONSUL_HOST', 'localhost'),
int(os.environ.get('CONSUL_PORT', 8500)),
os.environ.get('CONSUL_COUNTER_KEY', 'service/ids/counter'),
int(os.environ.get('CONSUL_MAX_RETRIES', 5)),
float(os.environ.get('CONSUL_RETRY_DELAY_SECONDS', 0.1))
)
+182 -85
View File
@@ -1,99 +1,196 @@
"""
OpenTelemetry configuration for tracing and metrics.
This module provides functionality to initialize and configure OpenTelemetry
for both tracing (Jaeger via OTLP) and metrics (Prometheus) in the AMQ adapter.
"""
import logging
import os
from dataclasses import dataclass
from typing import Any, Dict, Optional
from opentelemetry import metrics, trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.exporter.prometheus import PrometheusMetricReader
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
from opentelemetry.sdk.trace import Tracer, TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from opentelemetry.sdk.resources import SERVICE_NAME, SERVICE_VERSION, Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
from opentelemetry.trace import Tracer
from prometheus_client import start_http_server
logger = logging.getLogger(__name__)
def initialize_trace(app=None, name=None) -> Tracer:
# Resource can be required for some backends, e.g. Jaeger or Prometheus
# If resource wouldn't be set - traces wouldn't appear in Jaeger.
# The 'name' value is the name shown in Jaeger.
resource = Resource(
attributes={
SERVICE_NAME: (name if name else os.environ.get("APPLICATION_NAME", "Python App"))
}
)
#
# 1. Tracing (Jaeger) configuration
# ----------------------------------
#
# Configure the OTLP exporter (or HTTP exporter the same way)
# the recommended configuration is via environment variables. The mandatory ones:
# OTEL_EXPORTER_OTLP_ENDPOINT: "http://jaeger:4317"
# OTEL_EXPORTER_OTLP_INSECURE: "true"
# for more see here: https://opentelemetry-python.readthedocs.io/en/latest/exporter/otlp/otlp.html
# alternatively these values can be provided as input values to OTLPSpanExporter() constructor.
#
# otlp_exporter = OTLPSpanExporter()
# or provide endpoint explicitly, as
otlp_exporter = OTLPSpanExporter(endpoint="http://jaeger:4317")
#
# Initialize the tracer provider. This is abstract API call.
# The actual implementation is given by the included dependency
# The opentelemetry default is the OTLP trace provider.
@dataclass
class TracingConfig:
"""Configuration for OpenTelemetry tracing."""
service_name: str
service_version: str
otlp_endpoint: str = "http://localhost:4317" # Default Jaeger OTLP endpoint
insecure: bool = True
debug: bool = False
additional_attributes: Optional[Dict[str, Any]] = None
@dataclass
class MetricsConfig:
"""Configuration for OpenTelemetry metrics."""
prometheus_port: int = 9464
prometheus_host: str = "localhost"
export_interval_millis: int = 30000
additional_attributes: Optional[Dict[str, Any]] = None
def create_resource(config: TracingConfig) -> Resource:
"""Create an OpenTelemetry resource with service information."""
attributes = {
SERVICE_NAME: config.service_name,
SERVICE_VERSION: config.service_version,
}
if config.additional_attributes:
attributes.update(config.additional_attributes)
return Resource.create(attributes)
def setup_otlp_exporter(config: TracingConfig) -> OTLPSpanExporter:
"""Configure and create an OTLP exporter for Jaeger."""
return OTLPSpanExporter(endpoint=config.otlp_endpoint, insecure=config.insecure)
def setup_trace_provider(config: TracingConfig, resource: Resource) -> TracerProvider:
"""Set up the TracerProvider with OTLP exporter."""
provider = TracerProvider(resource=resource)
# Add the exporters to the batch span processor
span_processor = BatchSpanProcessor(otlp_exporter)
provider.add_span_processor(span_processor)
# Sets the global default tracer provider
trace.set_tracer_provider(provider)
#
# logging.info(' * Instrumenting application..... ')
# if Flask, then use:
# FlaskInstrumentor().instrument_app(flask_app)
# if FastAPI, then use
# FastAPIInstrumentor.instrument_app(app)
# Add OTLP exporter for Jaeger
otlp_exporter = setup_otlp_exporter(config)
provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
#
# Python Flask is now automatically instrumented to export details of each REST endpoint call.
# For Django based app, use 'opentelemetry-instrumentation-django' import instead. see here:
# https://opentelemetry-python.readthedocs.io/en/latest/examples/django/README.html for example
#
# If the used library has no instrumentation support in opentelemetry.io, then the trace
# can be started manually. In this code,
# 'trace' should be global shared variable, initialized once per application (see tracing.py)
# 'tracer' can be shared global or local variable, application can use multiple uniquely
# named tracers, Jaeger will group the generated traces by the name of the tracer
# that was used to start the Span. See this example code:
#
# tracer = trace.get_tracer(name)
#
# @app.route("/server_request")
# def server_request():
# with tracer.start_as_current_span(
# "server_request",
# context=extract(request.headers),
# kind=trace.SpanKind.SERVER,
# attributes=collect_request_attributes(request.environ),
# ):
# # do the work here, like logging.info something and return response
# logging.info(request.args.get("param"))
# return "served"
#
# Note:
# the otlp code looks for HTTP header named 'traceparent' that, if present, should contain
# the otlp parent span identification, and the span created here will be made a child
# of that parent span, indicating visually the progression of requests via different services
# for as long as the current span is provided in the traceparent header to next service.
# Add console exporter in debug mode
if config.debug:
provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))
# 2. Metrics (Prometheus) configuration
# -------------------------------------
#
# Start Prometheus client
start_http_server(port=9464, addr="localhost")
# Initialize PrometheusMetricReader which pulls metrics from the SDK
# on-demand to respond to scrape requests (unlike Jaeger, which worls with push model - application
# pushes traces when those become available to Jaeger's collector. Prometheus works with pull model
# where it periodically polls pre-configured and/or discovered endpoints for metrics.
reader = PrometheusMetricReader()
provider = MeterProvider(resource=resource, metric_readers=[reader])
metrics.set_meter_provider(provider)
return provider
return trace.get_tracer(name)
def setup_metrics_provider(config: MetricsConfig, resource: Resource) -> MeterProvider:
"""Set up the MeterProvider with Prometheus export."""
try:
# Start Prometheus HTTP server
start_http_server(port=config.prometheus_port, addr=config.prometheus_host)
logger.info(
f"Started Prometheus metrics server on {config.prometheus_host}:{config.prometheus_port}"
)
# Create Prometheus reader
prometheus_reader = PrometheusMetricReader()
# Create periodic reader for custom metrics
periodic_reader = PeriodicExportingMetricReader(
prometheus_reader, export_interval_millis=config.export_interval_millis
)
# Create and configure meter provider
return MeterProvider(resource=resource, metric_readers=[prometheus_reader, periodic_reader])
except Exception as e:
logger.error(f"Failed to set up metrics provider: {str(e)}")
# Return a no-op meter provider
return MeterProvider(resource=resource)
def initialize_telemetry(
tracing_config: Optional[TracingConfig] = None,
metrics_config: Optional[MetricsConfig] = None,
) -> Tracer:
"""
Initialize OpenTelemetry with tracing and metrics.
Args:
tracing_config: Configuration for tracing
metrics_config: Configuration for metrics
Returns:
Tracer: Configured OpenTelemetry tracer
Example:
```python
config = TracingConfig(
service_name="amq-adapter",
service_version="1.0.0",
otlp_endpoint="http://jaeger:4317",
debug=True
)
metrics_config = MetricsConfig(
prometheus_port=9464,
prometheus_host="0.0.0.0"
)
tracer = initialize_telemetry(config, metrics_config)
```
"""
# Use default configs if none provided
if tracing_config is None:
tracing_config = TracingConfig(
service_name=os.getenv("SERVICE_NAME", "amq-adapter"),
service_version=os.getenv("SERVICE_VERSION", "unknown"),
otlp_endpoint=os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317"),
insecure=os.getenv("OTEL_INSECURE_MODE", "true").lower() == "true",
debug=os.getenv("OTEL_DEBUG", "false").lower() == "true",
)
if metrics_config is None:
metrics_config = MetricsConfig(
prometheus_port=int(os.getenv("PROMETHEUS_PORT", "9464")),
prometheus_host=os.getenv("PROMETHEUS_HOST", "localhost"),
export_interval_millis=int(os.getenv("METRICS_EXPORT_INTERVAL", "30000")),
)
try:
# Create resource
resource = create_resource(tracing_config)
# Set up tracing
provider = setup_trace_provider(tracing_config, resource)
trace.set_tracer_provider(provider)
logger.info(f"Initialized tracing for service {tracing_config.service_name}")
# Set up metrics
metrics_provider = setup_metrics_provider(metrics_config, resource)
metrics.set_meter_provider(metrics_provider)
logger.info("Initialized metrics export")
# Return configured tracer
return trace.get_tracer(
tracing_config.service_name,
tracing_config.service_version,
)
except Exception as e:
logger.error(f"Failed to initialize telemetry: {str(e)}")
# Return a no-op tracer in case of failure
return trace.get_tracer("no-op")
def create_span(tracer: Tracer, name: str, attributes: Optional[Dict[str, Any]] = None):
"""
Create a new span with the given name and attributes.
Args:
tracer: The OpenTelemetry tracer
name: Name of the span
attributes: Optional attributes to add to the span
Returns:
A context manager that creates and manages a span
"""
return tracer.start_as_current_span(
name,
attributes=attributes or {},
)
+4 -1
View File
@@ -104,8 +104,11 @@ class CleverBragAmqpAdapter(CleverThisServiceAdapter):
amq_configuration: AMQConfiguration,
routes: List[APIRoute],
auth_provider: core.base.providers.auth.AuthProvider,
extra_init_data: dict,
):
super().__init__(amq_configuration, routes, None)
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)
+4 -1
View File
@@ -21,7 +21,10 @@ if __name__ == "__main__":
r2rapp: R2RApp = create_r2r_app()
adapter: CleverBragAmqpAdapter = CleverBragAmqpAdapter(
AMQConfiguration("./application.properties.local"), r2rapp.app.routes, r2rapp.providers.auth
AMQConfiguration("./application.properties.local"),
r2rapp.app.routes,
r2rapp.providers.auth,
{"app": r2rapp.config.app},
)
for method, routes in adapter.routes.items():
+149
View File
@@ -0,0 +1,149 @@
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
@@ -0,0 +1,42 @@
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.

After

Width:  |  Height:  |  Size: 201 KiB

Binary file not shown.
@@ -0,0 +1 @@
[{"@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
@@ -0,0 +1,180 @@
#!/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
@@ -0,0 +1,98 @@
#!/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
+337
View File
@@ -0,0 +1,337 @@
import json
import os
import random
import shutil
import time
import httpx # For emulating remote service call
# import uvicorn
from fastapi import FastAPI, File, Form, UploadFile
# OpenTelemetry imports
from opentelemetry import metrics, trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.exporter.prometheus import PrometheusMetricReader
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
from opentelemetry.sdk.metrics import MeterProvider
# from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
from amqp.config.amq_configuration import AMQConfiguration
from otdemo.otdemo_adapter import OTDemoAmqpAdapter
# --- OpenTelemetry Configuration ---
# Configure Resource: Defines common attributes for traces and metrics
resource = Resource.create(
attributes={
"service.name": "fastapi-document-service",
"application": "CleverMicroDemo",
"environment": "development",
}
)
# Configure TracerProvider: Responsible for creating and managing Tracers
trace_provider = TracerProvider(resource=resource)
# OTLPSpanExporter: Exports spans to an OTLP collector (e.g., Jaeger) via gRPC
# Endpoint can be overridden by OTEL_EXPORTER_OTLP_ENDPOINT environment variable
otlp_exporter = OTLPSpanExporter(
endpoint=os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317")
)
# BatchSpanProcessor: Batches spans for efficient export
trace_provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
# Optional: ConsoleSpanExporter for local debugging to print traces to console
if os.environ.get("OTEL_DEBUG_CONSOLE_TRACES", "false").lower() == "true":
trace_provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))
# Set the configured TracerProvider globally
trace.set_tracer_provider(trace_provider)
# Configure MeterProvider: Responsible for creating and managing Meters for metrics
# PrometheusMetricReader: Exposes metrics via an HTTP endpoint for Prometheus to scrape
# Host and port can be overridden by OTEL_EXPORTER_PROMETHEUS_HOST and OTEL_EXPORTER_PROMETHEUS_PORT
prometheus_reader = PrometheusMetricReader(
# handler_address=os.environ.get("OTEL_EXPORTER_PROMETHEUS_HOST", "0.0.0.0"),
# handler_port=int(os.environ.get("OTEL_EXPORTER_PROMETHEUS_PORT", "8001"))
)
# MeterProvider: Uses the PrometheusMetricReader to periodically export metrics
metric_provider = MeterProvider(resource=resource, metric_readers=[prometheus_reader])
# Set the configured MeterProvider globally
metrics.set_meter_provider(metric_provider)
# Get tracer and meter instances from the global providers
tracer = trace.get_tracer("document-service.tracer")
meter = metrics.get_meter("document-service.meter")
# --- Custom Metrics Definition ---
# Counters: For cumulative sums (e.g., total requests)
documents_uploaded_counter = meter.create_counter(
name="documents_uploaded_total", description="Total number of documents uploaded", unit="1"
)
metadata_processed_counter = meter.create_counter(
name="documents_metadata_processed_total",
description="Total number of document metadata processed",
unit="1",
)
documents_retrieved_counter = meter.create_counter(
name="documents_retrieved_total", description="Total number of documents retrieved", unit="1"
)
# Histograms: For distributions of values (e.g., request durations)
internal_api_call_duration_histogram = meter.create_histogram(
name="internal_api_call_duration_seconds",
description="Duration of internal API calls to metadata endpoint",
unit="s",
)
# --- FastAPI Application Setup ---
app = FastAPI(
title="CleverMicro Document Service Demo",
description="A FastAPI application demonstrating OpenTelemetry integration with Jaeger and Prometheus.",
# root_path="/otdemo", # This won't affect router path
)
# Instrument FastAPI: Automatically creates spans for incoming requests
FastAPIInstrumentor.instrument_app(app)
# Instrument httpx: Automatically propagates trace context for outgoing HTTP calls
HTTPXClientInstrumentor().instrument()
# Directory for storing uploaded files (temporary for demo)
UPLOAD_DIR = "/tmp/clevermicro_documents"
os.makedirs(UPLOAD_DIR, exist_ok=True)
@app.post("/api/v3/documents")
async def upload_document(
file: UploadFile = File(...), # Multipart file upload
name: str = Form(...), # Form field for document name
description: str = Form(...), # Form field for document description
):
# Create a custom span for this specific endpoint's logic
with tracer.start_as_current_span("upload_document_endpoint_logic"):
# Save the uploaded file to the temporary directory
file_location = os.path.join(UPLOAD_DIR, file.filename)
try:
with open(file_location, "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
print(f"File '{file.filename}' saved to {file_location}")
# Increment the documents_uploaded_total metric
documents_uploaded_counter.add(1, {"file.name": file.filename, "document.name": name})
# Emulate a REST call to a "remote" service (which is actually this same app)
# This demonstrates trace context propagation across HTTP calls.
metadata_payload = {"name": name, "description": description}
print(
f"Emulating remote call to /api/v3/documents/metadata with: {json.dumps(metadata_payload)}"
)
start_time = time.time()
async with httpx.AsyncClient() as client:
# httpx instrumentation ensures the current trace context is propagated
response = await client.post(
# IMPORTANT: Use the actual host and port where your FastAPI app is running
# If running directly on host, use http://localhost:8000
# If running inside Docker and calling itself, use http://host.docker.internal:8000
# "http://localhost:8000/api/v3/documents/metadata",
"http://app:8000/api/v3/documents/metadata",
json=metadata_payload,
)
response.raise_for_status() # Raise an exception for bad status codes (4xx or 5xx)
end_time = time.time()
duration = end_time - start_time
# Record the duration of the internal API call
internal_api_call_duration_histogram.record(
duration, {"endpoint": "/api/v3/documents/metadata"}
)
print(f"Internal metadata service responded: {response.json()}")
return {
"message": "Document uploaded and metadata processing initiated",
"filename": file.filename,
"name": name,
"description": description,
}
except httpx.HTTPStatusError as e:
print(f"Error calling metadata service: {e.response.status_code} - {e.response.text}")
# Propagate the error in the trace
current_span = trace.get_current_span()
current_span.set_attribute("error.type", "HTTPStatusError")
current_span.record_exception(e)
return {
"message": "Document uploaded but metadata processing failed",
"detail": str(e),
}, 500
except Exception as e:
print(f"An unexpected error occurred during document upload: {e}")
# Propagate the error in the trace
current_span = trace.get_current_span()
current_span.set_attribute("error.type", "UnexpectedError")
current_span.record_exception(e)
return {"message": "Error processing document", "detail": str(e)}, 500
@app.post("/api/v3/documents/metadata")
async def process_document_metadata(metadata: dict):
# This endpoint is automatically instrumented by FastAPIInstrumentor
# A new span will be created, and its parent will be the span from the calling /api/v3/documents
# due to trace context propagation via httpx.
print(f"Received metadata for processing: {json.dumps(metadata, indent=2)}")
# Increment the metadata_processed_total metric
metadata_processed_counter.add(1, {"document.name": metadata.get("name")})
# Simulate some processing time to make traces more interesting
time.sleep(random.uniform(0.05, 0.2))
return {"status": "metadata processed", "received_data": metadata}
@app.get("/api/v3/documents")
async def get_documents():
# Create a custom span for the logic within this endpoint
with tracer.start_as_current_span("get_documents_endpoint_logic"):
documents = []
# Generate several random documents to emulate a list
for i in range(random.randint(2, 5)):
documents.append(
{
"id": f"doc-{random.randint(1000, 9999)}",
"name": f"Document {random.randint(1, 100)}",
"description": f"Description for document {random.randint(1, 1000)}",
"uploaded_at": time.strftime(
"%Y-%m-%dT%H:%M:%SZ",
time.gmtime(time.time() - random.randint(0, 86400 * 30)),
),
}
)
# Increment the documents_retrieved_total metric
documents_retrieved_counter.add(1)
# Simulate some processing time
time.sleep(random.uniform(0.01, 0.1))
return {"documents": documents}
# Need to place below all route methods, otherwise the routes are incomplete
otadapter = OTDemoAmqpAdapter(AMQConfiguration("otdemo/otdemo.properties"), app.routes)
# --- How to Run the Demo ---
# 1. Save the code:
# Save the Python code above as `main.py` in a new directory.
# 2. Create `requirements.txt`:
# Create a file named `requirements.txt` in the same directory with the following content:
# ```
# fastapi==0.111.0
# uvicorn[standard]==0.29.0
# python-multipart==0.0.20
# httpx==0.27.0
# opentelemetry-api==1.25.0
# opentelemetry-sdk==1.25.0
# opentelemetry-exporter-otlp-proto-grpc==1.25.0 # For Jaeger traces
# opentelemetry-exporter-prometheus==1.25.0 # For Prometheus metrics
# opentelemetry-instrumentation-fastapi==0.46b0 # IMPORTANT: Check compatibility with your FastAPI version!
# opentelemetry-instrumentation-httpx==0.46b0 # IMPORTANT: Check compatibility with your httpx version!
# protobuf==4.25.3 # Required by OTLP exporter, ensure version compatibility if issues arise
# ```
# 3. Install dependencies:
# Open your terminal in the directory where you saved `main.py` and `requirements.txt`, then run:
# `pip install -r requirements.txt`
# 4. Run Jaeger and Prometheus (using Docker Compose for simplicity):
# Create a file named `docker-compose.yaml` in the same directory:
# ```yaml
# version: '3.8'
# services:
# jaeger:
# image: jaegertracing/all-in-one:latest
# ports:
# - "6831:6831/udp" # UDP Thrift
# - "14268:14268" # HTTP Thrift
# - "14250:14250" # gRPC
# - "4317:4317" # OTLP gRPC collector (for traces from our app)
# - "4318:4318" # OTLP HTTP collector
# - "16686:16686" # Jaeger UI
# environment:
# - COLLECTOR_OTLP_ENABLED=true # Enable OTLP reception
#
# prometheus:
# image: prom/prometheus:latest
# volumes:
# - ./prometheus.yml:/etc/prometheus/prometheus.yml # Mount Prometheus config
# ports:
# - "9090:9090" # Prometheus UI
# command:
# - '--config.file=/etc/prometheus/prometheus.yml'
# ```
# Create a file named `prometheus.yml` in the same directory:
# ```yaml
# global:
# scrape_interval: 15s # How frequently to scrape targets
#
# scrape_configs:
# - job_name: 'fastapi_app'
# # The 'metrics' endpoint of our FastAPI app exposed by PrometheusMetricReader
# # 'host.docker.internal' allows Docker containers to connect to the host machine's localhost
# static_configs:
# - targets: ['host.docker.internal:8001']
# ```
# Start Jaeger and Prometheus:
# `docker-compose up -d` (the `-d` runs them in the background)
# 5. Run the FastAPI application:
# In your terminal (where `main.py` is located), run:
# `uvicorn main:app --host 0.0.0.0 --port 8000`
# (Optional: To see traces printed to console, set `export OTEL_DEBUG_CONSOLE_TRACES=true` before `uvicorn`)
# --- Accessing the UIs ---
# * **Jaeger UI:** Open your web browser and go to `http://localhost:16686`
# * **Prometheus UI:** Open your web browser and go to `http://localhost:9090`
# * **FastAPI App (Root):** `http://localhost:8000`
# * **FastAPI Metrics Endpoint:** `http://localhost:8001/metrics` (Prometheus scrapes this, you can also view it directly)
# --- How to Test and Observe ---
# 1. Ensure all services are running (`docker-compose ps` should show `jaeger` and `prometheus` up, and `uvicorn` running in your terminal).
# 2. **Trigger GET requests:**
# Open your browser or use `curl` to hit: `http://localhost:8000/api/v3/documents`
# Refresh a few times.
# 3. **Trigger POST requests (for file upload and internal call):**
# You'll need a tool like Postman, Insomnia, or `curl` (more complex for multipart) for this.
# **Using Postman/Insomnia:**
# * Method: `POST`
# * URL: `http://localhost:8000/api/v3/documents`
# * Body: Select `form-data`
# * Add a key `file`, Type `File`, Value: Choose any small file from your computer (e.g., a `.txt` or `.png`).
# * Add a key `name`, Type `Text`, Value: `My Demo Document`
# * Add a key `description`, Type `Text`, Value: `A test document for OpenTelemetry.`
# * Send the request multiple times.
# 4. **Observe in Jaeger UI (`http://localhost:16686`):**
# * Select "Service": `fastapi-document-service`
# * Click "Find Traces".
# * You should see traces for `/api/v3/documents` (POST) and `/api/v3/documents` (GET).
# * Crucially, for POST requests, expand the trace: you will see the main `/api/v3/documents` span, and nested within it, a child span for the `POST /api/v3/documents/metadata` HTTP request. This demonstrates automatic context propagation.
# 5. **Observe in Prometheus UI (`http://localhost:9090`):**
# * Go to the "Graph" tab.
# * In the expression bar, type and execute queries like:
# * `documents_uploaded_total`
# * `documents_metadata_processed_total`
# * `documents_retrieved_total`
# * `internal_api_call_duration_seconds_sum`
# * `internal_api_call_duration_seconds_count`
# * You will see the values of these metrics increasing with your requests.
# This setup provides a clear, runnable demo for OpenTelemetry traces and metrics in a FastAPI application, including automatic context propagation for internal HTTP calls.
+65
View File
@@ -0,0 +1,65 @@
version: '3.8'
services:
app:
# build: .
image: git.cleverthis.com/clevermicro/amq-adapter-python-demo:latest
ports:
- "8000:8000"
- "8001:8001"
environment:
- OTEL_EXPORTER_OTLP_ENDPOINT=http://jaeger:4317
depends_on:
- jaeger
- prometheus
- rabbitmq
- alertmanager
jaeger:
image: jaegertracing/all-in-one:latest
ports:
- "6831:6831/udp"
- "14268:14268"
- "14250:14250"
- "4317:4317"
- "4318:4318"
- "16686:16686"
environment:
- COLLECTOR_OTLP_ENABLED=true
prometheus:
image: prom/prometheus:latest
volumes:
- ./container-data/prometheus/prometheus.yml:/etc/prometheus/prometheus.yml
- ./container-data/prometheus/alert-rules.yml:/etc/prometheus/alert-rules.yml
ports:
- "9090:9090"
command:
- '--config.file=/etc/prometheus/prometheus.yml'
depends_on:
- alertmanager
alertmanager:
image: prom/alertmanager:latest
ports:
- "9093:9093"
volumes:
- ./container-data/alertmanager/alertmanager.yml:/etc/alertmanager/alertmanager.yml
command:
- '--config.file=/etc/alertmanager/alertmanager.yml'
rabbitmq:
image: rabbitmq:3-management-alpine
ports:
- "5672:5672"
- "15672:15672"
- "15692:15692"
environment:
- RABBITMQ_DEFAULT_USER=guest
- RABBITMQ_DEFAULT_PASS=guest
volumes:
- rabbitmq_data:/var/lib/rabbitmq
labels:
prometheus-job: rabbitmq
volumes:
rabbitmq_data:
+44
View File
@@ -0,0 +1,44 @@
# ====================================================================
# 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=cleverswarm
# 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-dev route
cm.amq-adapter.route-mapping=otdemo:otdemo-exchange:otdemo-queue:api.cm.dev.cleverthis.com.test.python.#:5:0
# local route
#cm.amq-adapter.route-mapping=otdemo::otdemo:api.dev.localhost.otdemo.#:5:0
# Indicate if AMQ Adapter should respond with HTTP 401 Unauthorized if the request does not contain valid JWT token
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.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.
# 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-port=8080
# Backpressure monitor settings
# Define maximum number of messages of being processed in parallel before triggering backpressure OVERLOAD alert
cm.backpressure.threshold=1
# 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 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
+59
View File
@@ -0,0 +1,59 @@
"""
This module provides the CleverSwarm AMQP Adapter for handling API requests via AMQP.
The code IS MEANT to be part of CleverSwarm codebase, and has direct dependencies on the CleverSwarm codebase.
It is NOT MEANT to be used as a standalone module. It WILL NOT run outside the CleverSwarm endpoints context.
"""
from threading import Thread
from typing import Any, List
# CleverSwarm specific imports
from fastapi.routing import APIRoute
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter
from amqp.config.amq_configuration import AMQConfiguration
from amqp.service.amq_service import AMQService
PREFERRED_SUFFIX = "_json"
# =================================================================================================
# ======================== C l e v e r S w a r m A M Q A d a p t e r ========================
# =================================================================================================
class OTDemoAmqpAdapter(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.
"""
def __init__(self, amq_configuration: AMQConfiguration, routes: List[APIRoute]):
super().__init__(amq_configuration, routes=routes, session=None)
_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) -> str:
"""
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 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.
:param endpoint: the Callable for the endpoint
:return: new Callable representing the JSON wrapper endpoint, if it exists or the original Callable.
"""
return endpoint
+8 -3
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "amq_adapter"
version = "0.2.31"
version = "0.2.34"
description = "The Python implementation of the AMQ Adapter for CleverMicro"
readme = "README.md"
authors = [
@@ -29,7 +29,8 @@ dependencies = [
"opentelemetry-instrumentation-pika",
"pika",
"psutil",
"python_multipart"
"python-multipart",
"consul-kv"
]
[tool.setuptools.packages.find]
@@ -39,6 +40,9 @@ include = ["amqp*"]
[tool.black]
line-length = 100
[tool.pytest.ini_options]
asyncio_default_fixture_loop_scope = "module"
[project.optional-dependencies]
dev = [
"pytest>=7.0",
@@ -49,5 +53,6 @@ dev = [
"black",
"isort",
"flake8",
"pre-commit"
"pre-commit",
"pylint"
]
+12
View File
@@ -0,0 +1,12 @@
fastapi==0.111.0
uvicorn[standard]==0.29.0
python-multipart==0.0.20
httpx==0.27.0
opentelemetry-api==1.34.1
opentelemetry-sdk==1.34.1
opentelemetry-exporter-otlp-proto-grpc==1.34.1
opentelemetry-exporter-prometheus==0.55b1
opentelemetry-instrumentation-fastapi==0.55b1
opentelemetry-instrumentation-httpx==0.55b1
protobuf==5.28.3
consul-kv==0.7.4
View File
View File
@@ -0,0 +1,77 @@
import logging
import os
import threading
from amqp.service.consul_global_id_generator import get_unique_instance_id
CONSUL_HOST = os.environ.get("CONSUL_HOST", "localhost")
CONSUL_PORT = int(os.environ.get("CONSUL_PORT", 8500))
CONSUL_COUNTER_KEY = "service/ids/counter"
MAX_RETRIES = 5 # Max attempts for atomic update
RETRY_DELAY_SECONDS = 0.1 # Delay between retries for atomic update
# --- Logging Setup ---
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(threadName)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
# --- Demo: Simulate Concurrent Access ---
def worker_thread(results_list):
"""A function to be run by multiple threads to get unique IDs."""
instance_id = get_unique_instance_id()
results_list.append(instance_id)
logger.info(f"Thread {threading.current_thread().name} got ID: {instance_id}")
if __name__ == "__main__":
num_threads = 10 # Number of concurrent service instances to simulate
all_generated_ids = []
print(f"\n--- Simulating {num_threads} concurrent service instances requesting unique IDs ---")
threads = []
for i in range(num_threads):
thread = threading.Thread(
target=worker_thread, args=(all_generated_ids,), name=f"ServiceInstance-{i + 1}"
)
threads.append(thread)
thread.start()
# Wait for all threads to complete
for thread in threads:
thread.join()
print("\n--- All unique IDs generated ---")
print(f"Generated IDs: {sorted(all_generated_ids)}")
# Verify uniqueness
if len(all_generated_ids) == len(set(all_generated_ids)):
print("All generated IDs are unique (as expected).")
else:
print(
"WARNING: Duplicate IDs found! This might indicate a problem with atomicity or fallback logic."
)
# Identify duplicates
seen = set()
duplicates = set()
for x in all_generated_ids:
if x in seen:
duplicates.add(x)
seen.add(x)
print(f"Duplicate IDs: {duplicates}")
print("\n--- Testing Consul Unavailability Fallback ---")
# Temporarily change Consul port to simulate unavailability
original_consul_port = CONSUL_PORT
CONSUL_PORT = 1 # Use an unlikely port
# Clear any existing ConsulKV client cache if applicable (consul_kv typically creates new client per call)
# If using a persistent client, you'd need to invalidate/reinitialize it here.
fallback_id_1 = get_unique_instance_id()
fallback_id_2 = get_unique_instance_id()
print(f"Fallback ID 1: {fallback_id_1}")
print(f"Fallback ID 2: {fallback_id_2}")
+388
View File
@@ -0,0 +1,388 @@
#
# 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
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"
current_availability = 5
scale_request = ScaleRequestV1(
service_id,
task_id,
current_availability,
)
assert scale_request.version == 1
assert scale_request.serviceId == service_id
assert scale_request.taskId == task_id
assert scale_request.currentAvailability == current_availability
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"
current_availability = 5
scale_request = ScaleRequestV1(
service_id,
task_id,
current_availability,
)
json_output = scale_request.to_json()
expected_json = json.dumps(
{
"version": 1,
"serviceId": service_id,
"taskId": task_id,
"currentAvailability": current_availability,
},
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"
current_availability = 5
scale_request = ScaleRequestV1(
service_id,
task_id,
current_availability,
)
json_output = scale_request.to_json()
expected_json = json.dumps(
{
"version": 1,
"serviceId": service_id,
"taskId": task_id,
"currentAvailability": current_availability,
},
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
@pytest.mark.asyncio
async 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)
handler.send_backpressure_report = AsyncMock()
initial_executions = handler.current_parallel_executions
with patch("time.time") as mock_time:
mock_time.return_value = 12345
await handler.increase_parallel_executions()
assert handler.current_parallel_executions == initial_executions + 1
assert handler.last_usage_report_time == 12345
assert handler.send_backpressure_report.call_count == 1
@pytest.mark.asyncio
async 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.send_backpressure_report = AsyncMock()
handler.current_parallel_executions = 5
with patch("time.time") as mock_time:
mock_time.return_value = 12345
await handler.decrease_parallel_executions()
assert handler.current_parallel_executions == 4
assert handler.last_recover_report_time == 12345
assert handler.send_backpressure_report.call_count == 1
handler.last_recover_report_time = 0
handler.current_parallel_executions = 0
await handler.decrease_parallel_executions()
assert handler.current_parallel_executions == 0 # Should not go below 0
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_regular_report_loop(self, mock_channel, mock_loop, mock_config):
"""
Test the _regular_report_loop method.
"""
handler = BackpressureHandler(mock_channel, mock_loop, mock_config)
handler.send_backpressure_report = AsyncMock()
# test 1 loop
handler.do_loop = 1
await handler._regular_report_loop()
handler.send_backpressure_report.assert_called_once()
@pytest.mark.asyncio
async def test_send_backpressure_report(self, mock_channel, mock_loop, mock_config):
"""
Test the send_backpressure_report 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.send_backpressure_report()
assert handler.publish_backpressure_request.call_count == 1
@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", 50)
)
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", 50)
)
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()
+265 -59
View File
@@ -1,14 +1,105 @@
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
from amqp.adapter.cleverthis_service_adapter import (
AMQMessage,
CleverThisServiceAdapter,
_convert_file_response,
_create_reply_exchange_and_queue,
is_likely_json,
)
from amqp.adapter.file_handler import FileHandler
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import DataMessage, DataResponse
from amqp.model.model import DataResponse
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE
class TestCleverThisServiceAdapter(unittest.TestCase):
class 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
def setUp(self):
# Mock AMQConfiguration
self.mock_amq_config = AMQConfiguration("./application.properties.local")
@@ -20,9 +111,11 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
self.adapter = CleverThisServiceAdapter(
amq_configuration=self.mock_amq_config, routes=[], 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, "localhost")
self.assertEqual(self.adapter.service_host, "testhost")
self.assertEqual(self.adapter.service_port, 8080)
self.assertFalse(self.adapter.require_authenticated_user)
@@ -48,26 +141,27 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
@patch.object(CleverThisServiceAdapter, "get_current_user")
async def test_on_message_unauthenticated(self, mock_get_user):
mock_message = MagicMock(spec=AMQMessage)
mock_message.headers = {}
mock_message.path = "/some/path"
mock_message.method = "GET"
mock_message.id = "123"
mock_message.trace_info = {}
mock_message.headers.return_value = {}
mock_message.path.return_value = "/some/path"
mock_message.method.return_value = "GET"
mock_message.id.return_value = "123"
mock_message.trace_info.return_value = {}
# Test when authentication is required but not provided
self.adapter.require_authenticated_user = True
response = await self.adapter.on_message(mock_message)
self.assertEqual(response.response_code(), "401")
mock_get_user.assert_not_called()
@patch.object(CleverThisServiceAdapter, "get_current_user")
@patch.object(CleverThisServiceAdapter, "get")
async def test_on_message_authenticated(self, mock_get, mock_get_user):
async def test_on_message_authenticated(self, mock_get: MagicMock, mock_get_user):
mock_message = MagicMock(spec=AMQMessage)
mock_message.headers = {"Authorization": ["Bearer valid_token"]}
mock_message.path = "/some/path"
mock_message.method = "GET"
mock_message.id = "123"
mock_message.trace_info = {}
mock_message.headers.return_value = {"Authorization": ["Bearer valid_token"]}
mock_message.path.return_value = "/some/path"
mock_message.method.return_value = "GET"
mock_message.id.return_value = "123"
mock_message.trace_info.return_value = {}
mock_user = MagicMock()
mock_get_user.return_value = mock_user
@@ -75,6 +169,7 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
mock_response = MagicMock(spec=DataResponse)
mock_get.return_value = mock_response
self.adapter.require_authenticated_user = True
response = await self.adapter.on_message(mock_message)
self.assertEqual(response, mock_response)
@@ -84,11 +179,11 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
@patch.object(CleverThisServiceAdapter, "get_current_user")
async def test_on_message_login_path(self, mock_get_user):
mock_message = MagicMock(spec=AMQMessage)
mock_message.headers = {}
mock_message.path = "/login"
mock_message.method = "POST"
mock_message.id = "123"
mock_message.trace_info = {}
mock_message.headers.return_value = {}
mock_message.path.return_value = "/login"
mock_message.method.return_value = "POST"
mock_message.id.return_value = "123"
mock_message.trace_info.return_value = {}
mock_response = MagicMock(spec=DataResponse)
with patch.object(self.adapter, "post", return_value=mock_response) as mock_post:
@@ -108,11 +203,11 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
for method in methods:
with self.subTest(method=method):
mock_message = MagicMock(spec=AMQMessage)
mock_message.headers = {"Authorization": ["Bearer valid_token"]}
mock_message.path = "/some/path"
mock_message.method = method
mock_message.id = "123"
mock_message.trace_info = {}
mock_message.headers.return_value = {"Authorization": ["Bearer valid_token"]}
mock_message.path.return_value = "/some/path"
mock_message.method.return_value = method
mock_message.id.return_value = "123"
mock_message.trace_info.return_value = {}
with patch.object(
self.adapter, method.lower(), return_value=mock_response
@@ -122,11 +217,11 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
mock_method.assert_called_once()
async def test_get_with_session(self):
mock_message = MagicMock(spec=DataMessage)
mock_message.path = "/api/test"
mock_message.headers = {"X-Test": "value"}
mock_message.trace_info = {"trace-id": "123"}
mock_message.id = "msg123"
mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.path.return_value = "/api/test"
mock_message.headers.return_value = {"X-Test": "value"}
mock_message.trace_info.return_value = {"trace-id": "123"}
mock_message.id.return_value = "msg123"
mock_response = MagicMock(spec=ClientResponse)
mock_response.status = 200
@@ -146,26 +241,26 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
async def test_get_without_session(self):
self.adapter.session = None
mock_message = MagicMock(spec=DataMessage)
mock_message.path = "/api/test"
mock_message.id = "msg123"
mock_message.trace_info = {}
mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.path.return_value = "/api/test"
mock_message.id.return_value = "msg123"
mock_message.trace_info.return_value = {}
response = await self.adapter.get(mock_message, None)
self.assertEqual(response.response_code(), "404")
async def test_put_with_form_handling(self):
mock_message = MagicMock(spec=DataMessage)
mock_message.path = "/api/test"
mock_message.headers = {}
mock_message.trace_info = {}
mock_message.id = "msg123"
mock_message.body = b"test=value"
mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.path.return_value = "/api/test"
mock_message.headers.return_value = {}
mock_message.trace_info.return_value = {}
mock_message.id.return_value = "msg123"
mock_message.body.return_value = b"test=value"
mock_response = MagicMock(spec=ClientResponse)
mock_response.status = 200
mock_response.content_type = "application/json"
mock_response: DataResponse = MagicMock(spec=DataResponse)
mock_response.response_code.return_value = "200"
mock_response.content_type.return_value = "application/json"
mock_response.read = AsyncMock(return_value=b'{"result": "success"}')
self.mock_session.put.return_value = mock_response
@@ -181,15 +276,15 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
mock_handler.assert_called_once()
async def test_post_with_form_handling(self):
mock_message = MagicMock(spec=DataMessage)
mock_message.path = "/api/test"
mock_message.headers = {}
mock_message.trace_info = {}
mock_message.id = "msg123"
mock_message.body = b"test=value"
mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.path.return_value = "/api/test"
mock_message.headers.return_value = {}
mock_message.trace_info.return_value = {}
mock_message.id.return_value = "msg123"
mock_message.body.return_value = b"test=value"
mock_response = MagicMock(spec=ClientResponse)
mock_response.status = 201
mock_response: DataResponse = MagicMock(spec=DataResponse)
mock_response.response_code.return_value = "201"
mock_response.content_type = "application/json"
mock_response.read = AsyncMock(return_value=b'{"result": "created"}')
@@ -201,24 +296,135 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
wraps=self.adapter._authorize_forward_request,
) as mock_handler:
response = await self.adapter.post(mock_message, None)
self.assertEqual(response.response_code(), "201")
mock_handler.assert_called_once()
async def test_head_not_implemented(self):
mock_message = MagicMock(spec=DataMessage)
mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.id = "msg123"
with self.assertRaises(NotImplementedError):
await self.adapter.head(mock_message, None)
async def test_delete_not_implemented(self):
mock_message = MagicMock(spec=DataMessage)
mock_message.id = "msg123"
def test_basic_type_conversion(self):
async def test_method(x: int, y: float, z: bool, text: str):
pass
# Assuming delete is not implemented in base class
with self.assertRaises(NotImplementedError):
await self.adapter.delete(mock_message, None)
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")
if __name__ == "__main__":
+165
View File
@@ -0,0 +1,165 @@
import json
import unittest
from unittest.mock import MagicMock, patch
from amqp.adapter.data_parser import AMQDataParser, MultipartFormDataParser
class DummyAMQMessage:
def __init__(self, headers=None, body=b"", extra_data=None):
self._headers = headers or {"Content-Type": ["application/json"]}
self._body = body
self.extra_data = extra_data if extra_data is not None else {}
def headers(self):
return self._headers
def body(self):
return self._body
class TestAMQDataParser(unittest.TestCase):
def test_process_form_data(self):
cases = [
(
'{"form": {"foo": ["bar"]}, "files": {"file1": ["data"]}}',
{"foo": "bar", "file1": "data"},
),
({"form": {"a": ["1"]}, "files": {"b": ["2"]}}, {"a": "1", "b": "2"}),
({"form": {}, "files": {}}, {}),
]
for json_input, expected in cases:
result = AMQDataParser.process_form_data(json_input)
self.assertEqual(result, expected)
def test_parse_get_url(self):
url = "http://example.com/path?foo=1&bar=2&foo=3"
path, params = AMQDataParser.parse_get_url(url)
self.assertEqual(path, "/path")
self.assertTrue(params["foo"] == ["1", "3"] or params["foo"] == "1" or params["foo"] == "3")
self.assertEqual(params["bar"], "2")
def test_parse_request_body_json(self):
data = {"foo": "bar"}
msg = DummyAMQMessage(
headers={"Content-Type": ["application/json"]}, body=json.dumps(data).encode()
)
result = AMQDataParser.parse_request_body(msg)
self.assertEqual(result, data)
def test_parse_request_body_invalid_json(self):
msg = DummyAMQMessage(headers={"Content-Type": ["application/json"]}, body=b"notjson")
with patch("amqp.adapter.data_parser.logging_error") as log_err:
with self.assertRaises(ValueError):
AMQDataParser.parse_request_body(msg)
self.assertTrue(log_err.called)
def test_parse_request_body_urlencoded(self):
msg = DummyAMQMessage(
headers={"Content-Type": ["application/x-www-form-urlencoded"]}, body=b"foo=bar&baz=qux"
)
result = AMQDataParser.parse_request_body(msg)
self.assertEqual(result, {b"foo": b"bar", b"baz": b"qux"})
def test_parse_request_body_urlencoded_json(self):
data = {"form": {"foo": ["bar"]}, "files": {}}
msg = DummyAMQMessage(
headers={"Content-Type": ["application/x-www-form-urlencoded"]},
body=json.dumps(data).encode(),
)
result = AMQDataParser.parse_request_body(msg)
self.assertEqual(result, {"foo": "bar"})
def test_parse_request_body_invalid_urlencoded(self):
msg = DummyAMQMessage(
headers={"Content-Type": ["application/x-www-form-urlencoded"]}, body=b"not a json"
)
with patch("amqp.adapter.data_parser.logging_error") as log_err:
with self.assertRaises(ValueError):
AMQDataParser.parse_request_body(msg)
self.assertTrue(log_err.called)
def test_parse_request_body_text(self):
msg = DummyAMQMessage(headers={"Content-Type": ["text/plain"]}, body=b"hello world")
result = AMQDataParser.parse_request_body(msg)
self.assertEqual(result, {"text": b"hello world"})
def test_parse_request_body_xml(self):
xml = b"<root><foo>bar</foo><baz>qux</baz></root>"
msg = DummyAMQMessage(headers={"Content-Type": ["application/xml"]}, body=xml)
result = AMQDataParser.parse_request_body(msg)
self.assertEqual(result, {"foo": "bar", "baz": "qux"})
def test_parse_request_body_invalid_xml(self):
msg = DummyAMQMessage(
headers={"Content-Type": ["application/xml"]}, body=b"<root><foo></bar></root>"
)
with self.assertRaises(ValueError):
AMQDataParser.parse_request_body(msg)
def test_parse_request_body_unsupported(self):
msg = DummyAMQMessage(headers={"Content-Type": ["application/unknown"]}, body=b"")
with self.assertRaises(ValueError):
AMQDataParser.parse_request_body(msg)
class TestMultipartFormDataParser(unittest.TestCase):
def test_multipart_form_data_parser_json(self):
data = {"form": {"foo": ["bar"]}, "files": {}}
msg = DummyAMQMessage(
headers={"Content-Type": ["multipart/form-data"]},
body=json.dumps(data).encode(),
extra_data={"some": "data"},
)
parser = MultipartFormDataParser(msg)
self.assertTrue(parser.is_json)
self.assertEqual(parser.form_data, {"foo": "bar"})
self.assertTrue(parser.is_valid)
def test_multipart_form_data_parser_invalid_json(self):
msg = DummyAMQMessage(
headers={"Content-Type": ["multipart/form-data"]},
body=b"notjson",
extra_data={"some": "data"},
)
with patch("amqp.adapter.data_parser.logging_error") as log_err:
parser = MultipartFormDataParser(msg)
self.assertFalse(parser.is_valid)
self.assertTrue(log_err.called)
def test_multipart_form_data_parser_fallback_to_json(self):
msg = DummyAMQMessage(
headers={"Content-Type": ["multipart/form-data"]},
body=json.dumps({"form": {"foo": ["bar"]}, "files": {}}).encode(),
)
with patch(
"amqp.adapter.data_parser.python_multipart.parse_form", side_effect=Exception("fail")
):
parser = MultipartFormDataParser(msg)
self.assertTrue(parser.is_json)
self.assertEqual(parser.form_data, {"foo": "bar"})
self.assertTrue(parser.is_valid)
def test_multipart_form_data_parser_on_field_and_file(self):
msg = DummyAMQMessage(
headers={"Content-Type": ["multipart/form-data"]}, body=b"", extra_data={}
)
parser = MultipartFormDataParser(msg)
class DummyField:
field_name = b"foo"
value = b"bar"
class DummyFile:
field_name = b"file"
file_object = MagicMock()
file_object.seek = MagicMock()
size = 123
file_name = b"filename.txt"
parser.on_field(DummyField())
self.assertEqual(parser.form_data["foo"], b"bar")
parser.on_file(DummyFile())
self.assertIn("file", parser.form_data)
self.assertTrue(hasattr(parser.form_data["file"], "filename"))
self.assertEqual(parser.form_data["file"].filename, "filename.txt")
+263 -2
View File
@@ -1,11 +1,22 @@
import asyncio
import os
import tempfile
from unittest import TestCase
from asyncio import AbstractEventLoop
from collections import defaultdict
from pathlib import Path
from shutil import rmtree
from threading import Thread
from unittest import IsolatedAsyncioTestCase
from unittest.mock import AsyncMock, MagicMock, mock_open, patch
from aio_pika import IncomingMessage
from aio_pika.abc import AbstractRobustConnection
from amqp.adapter import file_handler
from amqp.adapter.file_handler import StreamingFileHandler
class TestFileHandler(TestCase):
class TestFileHandler(IsolatedAsyncioTestCase):
def test_ensure_directory_exists(self):
"""
Test function for ensure_directory_exists. Creates a temporary directory
@@ -48,3 +59,253 @@ class TestFileHandler(TestCase):
), f"Test Case 3 Version 3 Failed: {result3_v3} != {os.path.join(temp_dir, 'test3.3.txt')}"
print("All test cases passed!")
async def test_download_from_queue_to_file(self):
loop = asyncio.new_event_loop()
tempFile = tempfile.mktemp()
file_handler = StreamingFileHandler()
with patch.object(
StreamingFileHandler, "_get_connection", new_callable=AsyncMock
) as mock_connect:
mock_connection: AbstractRobustConnection = MagicMock(spec=AbstractRobustConnection)
mock_connect.return_value = mock_connection
mock_connection.channel = AsyncMock()
mock_channel = mock_connection.channel.return_value
mock_channel.get_queue = AsyncMock()
mock_queue = mock_channel.get_queue.return_value
mock_queue.iterator = MagicMock()
mock_queue_iter = mock_queue.iterator().__aenter__.return_value
mock_message: IncomingMessage = MagicMock(spec=IncomingMessage)
mock_message.body = b"test_data"
mock_message.body_size = len(b"test_data")
mock_message.headers = {"eof": 1}
mock_queue_iter.__aiter__.return_value = [mock_message]
# Act
result = await file_handler._download_from_queue_to_file(
"test_rabbitmq",
"test_queue",
filepath=tempFile,
loop=loop,
)
# clean up
os.remove(tempFile)
# Assert
assert result == (len(b"test_data"), 1)
mock_connect.assert_called_once_with("test_rabbitmq", loop=loop)
mock_channel.get_queue.assert_called_once_with("test_queue", ensure=False)
async def test_download_file_success(self):
# Arrange
file_handler = StreamingFileHandler()
loop = asyncio.new_event_loop()
loop_thread = Thread(target=loop.run_forever, daemon=True)
loop_thread.start()
rabbitmq_url = "amqp://test.rabbitmq"
file_info = {
"filename": "test_file.txt",
"size": 1024,
"streamName": "test_stream",
"queue_name": "test_queue",
"key": "test_key",
}
message_data = {}
output_dir = tempfile.mkdtemp()
file_handler.ensure_directory_exists = MagicMock(return_value=f"{output_dir}/test_file.txt")
file_handler._download_from_queue_to_file = AsyncMock(return_value=(123, 1))
# Act
result = await file_handler.download_file(
loop, rabbitmq_url, file_info, message_data, output_dir
)
# clean up
loop.stop()
rmtree(output_dir, ignore_errors=True)
# Assert
assert result == (123, 1)
file_handler.ensure_directory_exists.assert_called_once_with(f"{output_dir}/test_file.txt")
file_handler._download_from_queue_to_file.assert_called_once_with(
"amqp://test.rabbitmq", "test_queue", f"{output_dir}/test_file.txt", loop
)
async def test_download_file_queue_not_found(self):
# Arrange
file_handler = StreamingFileHandler()
loop = MagicMock(spec=AbstractEventLoop)
rabbitmq_url = "amqp://test.rabbitmq"
file_info = {
"filename": "test_file.txt",
"size": 1024,
"streamName": "test_stream",
"queue_name": None,
"key": "test_key",
}
message_data = {}
output_dir = "test_output"
# Act
result = await file_handler.download_file(
loop, rabbitmq_url, file_info, message_data, output_dir
)
# Assert
assert result == (0, 2)
async def test_download_from_queue_to_buffer(self):
loop = asyncio.new_event_loop()
file_handler = StreamingFileHandler()
with patch.object(
StreamingFileHandler, "_get_connection", new_callable=AsyncMock
) as mock_connect:
mock_connection: AbstractRobustConnection = MagicMock(spec=AbstractRobustConnection)
mock_connect.return_value = mock_connection
mock_connection.channel = AsyncMock()
mock_channel = mock_connection.channel.return_value
mock_channel.get_queue = AsyncMock()
mock_queue = mock_channel.get_queue.return_value
mock_queue.iterator = MagicMock()
mock_queue_iter = mock_queue.iterator().__aenter__.return_value
mock_message: IncomingMessage = MagicMock(spec=IncomingMessage)
mock_message.body = b"test_data"
mock_message.body_size = len(b"test_data")
mock_message.headers = {"eof": 1}
mock_queue_iter.__aiter__.return_value = [mock_message]
# Act
result = await file_handler._download_from_queue_to_buffer(
"test_rabbitmq",
"test_queue",
loop=loop,
)
# Assert
assert result == (b"test_data", 1)
mock_connect.assert_called_once_with("test_rabbitmq", loop=loop)
mock_channel.get_queue.assert_called_once_with("test_queue", ensure=False)
async def test_download_buffer_success(self):
# Arrange
file_handler = StreamingFileHandler()
loop = asyncio.new_event_loop()
loop_thread = Thread(target=loop.run_forever, daemon=True)
loop_thread.start()
rabbitmq_url = "amqp://test.rabbitmq"
file_handler._download_from_queue_to_buffer = AsyncMock(return_value=(b"hello", 1))
# Act
result = await file_handler.download_buffer(loop, rabbitmq_url, "test_queue")
# clean up
loop.stop()
# Assert
assert result == (b"hello", 1)
file_handler._download_from_queue_to_buffer.assert_called_once_with(
"amqp://test.rabbitmq", "test_queue", loop
)
async def test_download_buffer_queue_not_found(self):
# Arrange
file_handler = StreamingFileHandler()
loop = MagicMock(spec=AbstractEventLoop)
rabbitmq_url = "amqp://test.rabbitmq"
# Act
result = await file_handler.download_buffer(loop, rabbitmq_url, None)
# Assert
assert result == (0, 2)
async def test_on_message_no_files_to_download(self):
handler = StreamingFileHandler()
message = MagicMock()
message.body = b"{}"
json_data = b'{"files": {}}'
rabbitmq_url = "amqp://localhost"
loop = AsyncMock()
output_dir = "downloaded_files"
handler.download_file = AsyncMock()
handler.download_file.return_value = (123, 1)
with patch("amqp.adapter.file_handler.logging_info"):
result = await handler.on_message(message, json_data, rabbitmq_url, loop, output_dir)
assert isinstance(result, defaultdict)
self.assertEqual(result, defaultdict())
async def test_on_message_with_files_to_download(self):
handler = StreamingFileHandler()
message = MagicMock()
message.body = b'{"fileStream1": "queue1"}'
json_data = b'{"files": {"key1": [{"streamName": "fileStream1"}]}}'
rabbitmq_url = "amqp://localhost"
loop = AsyncMock()
output_dir = "downloaded_files"
handler.download_file = AsyncMock()
with patch("amqp.adapter.file_handler.logging_info") as mock_logging_info:
result = await handler.on_message(message, json_data, rabbitmq_url, loop, output_dir)
self.assertEqual(result, file_handler.download_locations)
handler.download_file.assert_awaited_once()
mock_logging_info.assert_any_call(
f" [{message.delivery_tag}] about to create task for {{'streamName': 'fileStream1', 'queue_name': 'queue1', 'key': 'key1'}}"
)
mock_logging_info.assert_called_with(f"All files downloaded. d-tag: {message.delivery_tag}")
async def test_on_message_handles_exception(self):
handler = StreamingFileHandler()
message = MagicMock()
message.body = b"invalid_json"
message.ack = AsyncMock()
json_data = b"invalid_json"
rabbitmq_url = "amqp://localhost"
loop = AsyncMock()
output_dir = "downloaded_files"
with patch("amqp.adapter.file_handler.logging_error") as mock_logging_error:
await handler.on_message(message, json_data, rabbitmq_url, loop, output_dir)
mock_logging_error.assert_called()
async def test_publish_file_file_not_found(self):
file_handler = StreamingFileHandler()
loop = asyncio.get_event_loop()
exchange_mock = AsyncMock()
file_path = Path("/path/to/nonexistent_file.txt")
with self.assertRaises(FileNotFoundError):
await file_handler.publish_file(
exchange=exchange_mock,
file_path=file_path,
routing_key="test_routing_key",
loop=loop,
)
async def test_publish_file_success(self):
file_handler = StreamingFileHandler()
loop = asyncio.get_event_loop()
exchange_mock = AsyncMock()
file_path = Path("/path/to/test_file.txt")
data = b"test-data"
with (
patch("amqp.adapter.file_handler.Path.is_file", return_value=True),
patch("amqp.adapter.file_handler.os.path.getsize", return_value=len(data)),
patch("builtins.open", new=mock_open(read_data=data)) as mock_file,
):
result = await file_handler.publish_file(
exchange=exchange_mock,
file_path=file_path,
routing_key="test_routing_key",
loop=loop,
max_chunk_size=512,
)
mock_file.assert_called_with(file_path, "rb")
self.assertEqual(result, ("test_routing_key", len(data)))
+138
View File
@@ -0,0 +1,138 @@
import os
from importlib.metadata import version
from unittest.mock import patch
from amqp.adapter.logging_utils import (
get_context_info,
logging_debug,
logging_error,
logging_info,
logging_warning,
)
__version__ = version("amq_adapter") # Name must match pyproject.toml [project] name
__verbose__ = os.environ.get("AMQ_VERBOSE_LOGGING", "0") == "1"
class TestLoggingFunctions:
"""
A class containing pytest unit tests for the logging functions.
"""
@patch("logging.error")
def test_logging_error_with_exception(self, mock_logger):
"""
Test logging_error function when called within an exception handler.
"""
try:
raise ValueError("Test error")
except ValueError as ve:
logging_error("An error occurred: %s", ve)
mock_logger.assert_called_once()
# Assert that the traceback information is included in the log message
assert len(mock_logger.call_args_list) == 1
_val = mock_logger.call_args_list[0][0][0]
assert "[E] 'An error occurred: %s'" in _val
assert "test_logging_error_with_exception" in _val
@patch("logging.error")
def test_logging_error_without_exception(self, mock_logger):
"""
Test logging_error function when called without an active exception.
"""
logging_error("A regular error")
mock_logger.assert_called_once()
assert "A regular error" in mock_logger.call_args_list[0][0][0]
@patch("logging.warning")
def test_logging_warning_verbose_on(self, mock_logger):
"""
Test logging_warning function when verbose logging is enabled.
"""
with patch(
"amqp.adapter.logging_utils.__verbose__", True
): # Simulate verbose logging
logging_warning("A warning message")
mock_logger.assert_called_once()
assert "A warning message" in mock_logger.call_args_list[0][0][0]
@patch("logging.warning")
def test_logging_warning_verbose_off(self, mock_logger):
"""
Test logging_warning function when verbose logging is disabled.
"""
with patch(
"amqp.adapter.logging_utils.__verbose__", False
): # Simulate non-verbose logging
logging_warning("A warning message")
mock_logger.assert_called_once()
assert "A warning message" in mock_logger.call_args_list[0][0][0]
@patch("logging.info")
def test_logging_info_verbose_on(self, mock_logger):
"""
Test logging_info function with verbose logging enabled.
"""
with patch("amqp.adapter.logging_utils.__verbose__", True):
logging_info("An info message")
mock_logger.assert_called_once()
assert "An info message" in mock_logger.call_args_list[0][0][0]
@patch("logging.info")
def test_logging_info_verbose_off(self, mock_logger):
"""
Test logging_info function with verbose logging disabled.
"""
with patch("amqp.adapter.logging_utils.__verbose__", False):
logging_info("An info message")
mock_logger.assert_called_once()
assert "An info message" in mock_logger.call_args_list[0][0][0]
@patch("logging.debug")
def test_logging_debug(self, mock_logger):
"""
Test logging_debug function.
"""
logging_debug("A debug message")
mock_logger.assert_called_once()
assert "A debug message" in mock_logger.call_args_list[0][0][0]
def test_get_context_info_normal(self):
"""
Test get_context_info function in a normal function call.
"""
def dummy_function():
return get_context_info()
module_name, line_number, function_name = dummy_function()
assert (
"test_get_context_info_normal" in function_name
) # changed from co_name to function name
assert (
"test_logging_utils.py" in module_name
) # changed from None to test_unit.py
assert isinstance(line_number, int)
def test_get_context_info_no_frame(self):
"""
Test get_context_info function when no frame is available.
"""
# Simulate a situation where inspect.currentframe() returns None (which is hard to do directly)
with patch("inspect.currentframe") as mock_frame:
mock_frame.return_value = None
module_name, line_number, function_name = get_context_info()
assert module_name is None
assert line_number is None
assert function_name is None
def test_get_context_info_exception(self):
"""
Test get_context_info function when an exception occurs.
"""
with patch("inspect.currentframe") as mock_frame:
mock_frame.side_effect = Exception("Simulated exception")
module_name, line_number, function_name = get_context_info()
assert module_name is None
assert line_number is None
assert function_name is None
+205 -17
View File
@@ -1,38 +1,226 @@
import json
from datetime import date, datetime
from unittest import TestCase
from uuid import UUID
from amqp.adapter.pydantic_serializer import parse_url_query
from pydantic import BaseModel
from amqp.adapter.pydantic_serializer import (
deserialize_object,
parse_url_query,
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 = []
class PydanticSerializerTest(TestCase):
def setUp(self):
"""
Prepare the test data.
"""
self.test_uuid = UUID("12345678-1234-5678-1234-567812345678")
self.test_datetime = datetime(2024, 1, 1, 12, 0, 0)
self.nested_model = NestedModel(value="test", number=42)
self.test_model = TestModel(
name="test",
age=25,
uuid_field=self.test_uuid,
datetime_field=self.test_datetime,
nested_model=self.nested_model,
list_field=[1, 2, 3],
)
def test_serialize_object(self):
assert True
serialized = serialize_object(self.test_model)
self.assertIsInstance(serialized, str)
self.assertIn("TestModel", serialized)
self.assertIn("test", serialized)
self.assertIn("25", serialized)
self.assertIn(str(self.test_uuid), serialized)
self.assertIn(self.test_datetime.isoformat(), serialized)
def test_deserialize_object(self):
assert True
serialized = serialize_object(self.test_model)
deserialized = deserialize_object(serialized)
def test__convert_value(self):
assert True
self.assertIsInstance(deserialized, TestModel)
self.assertEqual(deserialized.name, "test")
self.assertEqual(deserialized.age, 25)
self.assertEqual(deserialized.uuid_field, self.test_uuid)
self.assertEqual(deserialized.datetime_field, self.test_datetime)
self.assertIsInstance(deserialized.nested_model, NestedModel)
self.assertEqual(deserialized.nested_model.value, "test")
self.assertEqual(deserialized.nested_model.number, 42)
self.assertEqual(deserialized.list_field, [1, 2, 3])
def test__is_uuid_string(self):
assert True
def test_serialize_deserialize_with_none(self):
model = TestModel(
name="test",
age=25,
uuid_field=self.test_uuid,
datetime_field=self.test_datetime,
nested_model=None,
list_field=[],
)
serialized = serialize_object(model)
deserialized = deserialize_object(serialized)
self.assertIsNone(deserialized.nested_model)
def test__is_datetime_string(self):
assert True
def test_serialize_deserialize_with_list_of_models(self):
nested_models = [NestedModel(value=f"test{i}", number=i) for i in range(3)]
model = TestModel(
name="test",
age=25,
uuid_field=self.test_uuid,
datetime_field=self.test_datetime,
nested_model=None,
list_field=nested_models,
)
serialized = serialize_object(model)
deserialized = deserialize_object(serialized)
self.assertEqual(len(deserialized.list_field), 3)
for i, item in enumerate(deserialized.list_field):
self.assertIsInstance(item, NestedModel)
self.assertEqual(item.value, f"test{i}")
self.assertEqual(item.number, i)
def test__is_int_string(self):
assert True
def test_parse_url_query_no_params(self):
test_url = "http://example.com/path"
path, params = parse_url_query(test_url)
def test_serialize_to_json(self):
assert True
self.assertEqual(path, "/path")
self.assertEqual(params, {})
def test_parse_url_query_single_value(self):
test_url = "http://example.com/path?foo=1"
path, params = parse_url_query(test_url)
self.assertEqual(path, "/path")
self.assertEqual(params["foo"], "1")
def test_parse_url_query(self):
path, args = parse_url_query(
"http://localhost:8080/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741?response_type=JsonFile"
)
assert path == "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741"
assert args == {"response_type": "JsonFile"}
self.assertEqual(path, "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741")
self.assertEqual(args, {"response_type": "JsonFile"})
path, args = parse_url_query(
"http://localhost:8080/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741"
)
assert path == "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741"
assert args == {}
self.assertEqual(path, "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741")
self.assertEqual(args, {})
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
+114
View File
@@ -0,0 +1,114 @@
import unittest
from io import BytesIO
from amqp.adapter.serializer import (
bytes_to_long,
csv_as_list,
long_to_bytes,
map_as_string,
map_with_list_as_string,
object_or_list_as_string,
parse_map_list_string,
parse_map_string,
read_long,
str_to_bool,
)
class TestSerializer(unittest.TestCase):
def test_long_to_bytes_and_bytes_to_long(self):
# Test conversion in both directions
test_numbers = [
-9223372036854775808,
-1,
0,
1,
9223372036854775807,
] # Min and max long values
for num in test_numbers:
bytes_val = long_to_bytes(num)
self.assertEqual(bytes_to_long(bytes_val), num)
def test_read_long(self):
# Test reading valid long
test_num = 12345
bio = BytesIO(long_to_bytes(test_num))
self.assertEqual(read_long(bio), test_num)
# Test reading from empty BytesIO
empty_bio = BytesIO()
self.assertEqual(read_long(empty_bio), 0)
def test_object_or_list_as_string(self):
# Test with string
self.assertEqual(object_or_list_as_string("test"), "test")
# Test with list
self.assertEqual(object_or_list_as_string(["a", "b", "c"]), "[a,b,c]")
# Test with empty list
self.assertEqual(object_or_list_as_string([]), "[]")
def test_map_as_string(self):
# Test with simple key-value pairs
test_map = {"key1": "value1", "key2": "value2"}
self.assertEqual(map_as_string(test_map), "key1=value1,key2=value2")
# Test with empty map
self.assertEqual(map_as_string({}), "")
def test_map_with_list_as_string(self):
# Test with lists as values
test_map = {"key1": ["a", "b"], "key2": ["c", "d"]}
self.assertEqual(map_with_list_as_string(test_map), "key1=[a,b],key2=[c,d]")
# Test with empty map
self.assertEqual(map_with_list_as_string({}), "")
def test_parse_map_string(self):
# Test valid map string
test_str = "{key1=value1,key2=value2}"
expected = {"key1": "value1", "key2": "value2"}
self.assertEqual(parse_map_string(test_str), expected)
# Test empty map string
self.assertEqual(parse_map_string("{}"), {})
# Test invalid format
self.assertEqual(parse_map_string("invalid"), {})
def test_csv_as_list(self):
# Test normal CSV string
self.assertEqual(csv_as_list("a,b,c"), ["a", "b", "c"])
# Test with spaces
self.assertEqual(csv_as_list(" a , b , c "), ["a", "b", "c"])
# Test with brackets
self.assertEqual(csv_as_list("[a,b,c]"), ["a", "b", "c"])
# Test empty list
self.assertEqual(csv_as_list("[]"), [""])
def test_parse_map_list_string(self):
# Test valid map list string
test_str = "{key1=[a,b],key2=[c,d]}"
expected = {"key1": ["a", "b"], "key2": ["c", "d"]}
self.assertEqual(parse_map_list_string(test_str), expected)
# Test empty map
self.assertEqual(parse_map_list_string("{}"), {})
# Test single entry
self.assertEqual(parse_map_list_string("{key1=[a,b]}"), {"key1": ["a", "b"]})
def test_str_to_bool(self):
# Test true values
true_values = ["true", "yes", "1", "t", "y", "on", "TRUE", "YES", "T", "Y", "ON"]
for value in true_values:
self.assertTrue(str_to_bool(value))
# Test false values
false_values = ["false", "no", "0", "f", "n", "off", "FALSE", "NO", "F", "N", "OFF"]
for value in false_values:
self.assertFalse(str_to_bool(value))
@@ -1,9 +1,9 @@
from unittest import TestCase
from amqp.adapter.service_message_factory import (
ServiceMessageFactory,
RS,
INVALID_MESSAGE,
RS,
ServiceMessageFactory,
)
from amqp.model.model import ServiceMessage
from amqp.model.service_message_type import ServiceMessageType
@@ -13,15 +13,11 @@ 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)
@@ -57,3 +53,4 @@ 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
@@ -0,0 +1,251 @@
"""Unit tests for the tracing module."""
import unittest
from unittest.mock import MagicMock, patch
from opentelemetry.sdk.resources import Resource
from opentelemetry.trace import Tracer
from amqp.service.tracing import (
MetricsConfig,
TracingConfig,
create_resource,
create_span,
initialize_telemetry,
setup_metrics_provider,
setup_otlp_exporter,
setup_trace_provider,
)
class TestTracingConfig(unittest.TestCase):
"""Test cases for TracingConfig."""
def test_default_values(self):
"""Test that default values are set correctly."""
config = TracingConfig(service_name="test-service", service_version="1.0.0")
self.assertEqual(config.service_name, "test-service")
self.assertEqual(config.service_version, "1.0.0")
self.assertEqual(config.otlp_endpoint, "http://localhost:4317")
self.assertTrue(config.insecure)
self.assertFalse(config.debug)
self.assertIsNone(config.additional_attributes)
def test_custom_values(self):
"""Test that custom values are set correctly."""
config = TracingConfig(
service_name="test-service",
service_version="1.0.0",
otlp_endpoint="http://custom:4317",
insecure=False,
debug=True,
additional_attributes={"custom": "attribute"},
)
self.assertEqual(config.otlp_endpoint, "http://custom:4317")
self.assertFalse(config.insecure)
self.assertTrue(config.debug)
self.assertEqual(config.additional_attributes, {"custom": "attribute"})
class TestMetricsConfig(unittest.TestCase):
"""Test cases for MetricsConfig."""
def test_default_values(self):
"""Test that default values are set correctly."""
config = MetricsConfig()
self.assertEqual(config.prometheus_port, 9464)
self.assertEqual(config.prometheus_host, "localhost")
self.assertEqual(config.export_interval_millis, 30000)
self.assertIsNone(config.additional_attributes)
def test_custom_values(self):
"""Test that custom values are set correctly."""
config = MetricsConfig(
prometheus_port=8000,
prometheus_host="0.0.0.0",
export_interval_millis=15000,
additional_attributes={"custom": "attribute"},
)
self.assertEqual(config.prometheus_port, 8000)
self.assertEqual(config.prometheus_host, "0.0.0.0")
self.assertEqual(config.export_interval_millis, 15000)
self.assertEqual(config.additional_attributes, {"custom": "attribute"})
class TestCreateResource(unittest.TestCase):
"""Test cases for create_resource function."""
def test_create_resource_basic(self):
"""Test creating a resource with basic configuration."""
config = TracingConfig(service_name="test-service", service_version="1.0.0")
resource = create_resource(config)
self.assertIsInstance(resource, Resource)
self.assertEqual(resource.attributes.get("service.name"), "test-service")
self.assertEqual(resource.attributes.get("service.version"), "1.0.0")
def test_create_resource_with_attributes(self):
"""Test creating a resource with additional attributes."""
config = TracingConfig(
service_name="test-service",
service_version="1.0.0",
additional_attributes={"custom": "attribute"},
)
resource = create_resource(config)
self.assertEqual(resource.attributes.get("custom"), "attribute")
class TestSetupExporters(unittest.TestCase):
"""Test cases for exporter setup functions."""
@patch("amqp.service.tracing.OTLPSpanExporter")
def test_setup_otlp_exporter(self, mock_otlp):
"""Test OTLP exporter setup."""
config = TracingConfig(
service_name="test-service",
service_version="1.0.0",
otlp_endpoint="http://jaeger:4317",
insecure=True,
)
setup_otlp_exporter(config)
mock_otlp.assert_called_once_with(endpoint="http://jaeger:4317", insecure=True)
@patch("amqp.service.tracing.OTLPSpanExporter")
def test_setup_otlp_exporter_secure(self, mock_otlp):
"""Test OTLP exporter setup with secure connection."""
config = TracingConfig(
service_name="test-service",
service_version="1.0.0",
otlp_endpoint="https://jaeger:4317",
insecure=False,
)
setup_otlp_exporter(config)
mock_otlp.assert_called_once_with(endpoint="https://jaeger:4317", insecure=False)
class TestSetupTraceProvider(unittest.TestCase):
"""Test cases for trace provider setup."""
@patch("amqp.service.tracing.TracerProvider")
@patch("amqp.service.tracing.BatchSpanProcessor")
@patch("amqp.service.tracing.setup_otlp_exporter")
def test_setup_trace_provider_basic(self, mock_otlp_exporter, mock_processor, mock_provider):
"""Test basic trace provider setup."""
config = TracingConfig(service_name="test-service", service_version="1.0.0")
resource = create_resource(config)
mock_provider_instance = MagicMock()
mock_provider.return_value = mock_provider_instance
setup_trace_provider(config, resource)
mock_provider.assert_called_once_with(resource=resource)
mock_otlp_exporter.assert_called_once()
self.assertEqual(mock_provider_instance.add_span_processor.call_count, 1)
@patch("amqp.service.tracing.TracerProvider")
@patch("amqp.service.tracing.BatchSpanProcessor")
@patch("amqp.service.tracing.setup_otlp_exporter")
@patch("amqp.service.tracing.ConsoleSpanExporter")
def test_setup_trace_provider_with_debug(
self, mock_console_exporter, mock_otlp_exporter, mock_processor, mock_provider
):
"""Test trace provider setup with debug mode."""
config = TracingConfig(service_name="test-service", service_version="1.0.0", debug=True)
resource = create_resource(config)
mock_provider_instance = MagicMock()
mock_provider.return_value = mock_provider_instance
setup_trace_provider(config, resource)
self.assertEqual(mock_provider_instance.add_span_processor.call_count, 2)
mock_console_exporter.assert_called_once()
class TestSetupMetricsProvider(unittest.TestCase):
"""Test cases for metrics provider setup."""
@patch("amqp.service.tracing.start_http_server")
@patch("amqp.service.tracing.PrometheusMetricReader")
@patch("amqp.service.tracing.PeriodicExportingMetricReader")
@patch("amqp.service.tracing.MeterProvider")
def test_setup_metrics_provider_success(
self, mock_meter_provider, mock_periodic_reader, mock_prometheus_reader, mock_start_server
):
"""Test successful metrics provider setup."""
config = MetricsConfig(prometheus_port=9464, prometheus_host="localhost")
resource = Resource({})
setup_metrics_provider(config, resource)
mock_start_server.assert_called_once_with(port=9464, addr="localhost")
mock_prometheus_reader.assert_called_once()
mock_periodic_reader.assert_called_once()
mock_meter_provider.assert_called_once()
@patch("amqp.service.tracing.start_http_server")
def test_setup_metrics_provider_failure(self, mock_start_server):
"""Test metrics provider setup with failure."""
mock_start_server.side_effect = Exception("Failed to start server")
config = MetricsConfig()
resource = Resource({})
provider = setup_metrics_provider(config, resource)
self.assertIsNotNone(provider) # Should return no-op provider
class TestInitializeTelemetry(unittest.TestCase):
"""Test cases for initialize_telemetry function."""
@patch("amqp.service.tracing.setup_trace_provider")
@patch("amqp.service.tracing.setup_metrics_provider")
@patch("amqp.service.tracing.trace")
@patch("amqp.service.tracing.metrics")
def test_initialize_telemetry_success(
self, mock_metrics, mock_trace, mock_setup_metrics, mock_setup_trace
):
"""Test successful telemetry initialization."""
config = TracingConfig(service_name="test-service", service_version="1.0.0")
metrics_config = MetricsConfig()
mock_tracer = MagicMock(spec=Tracer)
mock_trace.get_tracer.return_value = mock_tracer
tracer = initialize_telemetry(config, metrics_config)
self.assertEqual(tracer, mock_tracer)
mock_setup_trace.assert_called_once()
mock_setup_metrics.assert_called_once()
@patch("amqp.service.tracing.setup_trace_provider")
def test_initialize_telemetry_failure(self, mock_setup_trace):
"""Test telemetry initialization with failure."""
mock_setup_trace.side_effect = Exception("Setup failed")
tracer = initialize_telemetry()
self.assertIsNotNone(tracer) # Should return no-op tracer
class TestCreateSpan(unittest.TestCase):
"""Test cases for create_span function."""
def test_create_span_basic(self):
"""Test basic span creation."""
mock_tracer = MagicMock(spec=Tracer)
create_span(mock_tracer, "test-span")
mock_tracer.start_as_current_span.assert_called_once_with("test-span", attributes={})
def test_create_span_with_attributes(self):
"""Test span creation with attributes."""
mock_tracer = MagicMock(spec=Tracer)
attributes = {"custom": "attribute"}
create_span(mock_tracer, "test-span", attributes)
mock_tracer.start_as_current_span.assert_called_once_with(
"test-span", attributes=attributes
)
if __name__ == "__main__":
unittest.main()
+211
View File
@@ -0,0 +1,211 @@
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()
+211
View File
@@ -0,0 +1,211 @@
import base64
import json
import unittest
from amqp.model.model import (
AMQErrorMessage,
AMQRoute,
CleverMicroMessage,
DataMessage,
DataResponse,
ScalingRequest,
)
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",
current_availability=3,
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["currentAvailability"], 3)
self.assertEqual(data["version"], 1)
def test_from_json(self):
json_str = json.dumps(
{
"serviceId": "test_service",
"taskId": "test_task",
"currentAvailability": 3,
"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.current_availability, 3)
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
@@ -0,0 +1,209 @@
from unittest import IsolatedAsyncioTestCase
from unittest.mock import AsyncMock, MagicMock, patch
from amqp.model.model import DataMessage
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient
class TestRabbitMQClient(IsolatedAsyncioTestCase):
def setUp(self):
self.mock_config = MagicMock()
self.mock_config.dispatch.amq_host = "localhost"
self.mock_config.dispatch.amq_port = 5672
self.mock_config.dispatch.rabbit_mq_user = "user"
self.mock_config.dispatch.rabbit_mq_password = "pass"
self.mock_config.dispatch.use_dlq = False
patcher = patch("amqp.router.router_consumer.RouterConsumer")
self.addCleanup(patcher.stop)
self.mock_router = patcher.start()
self.client = RabbitMQClient(self.mock_config, self.mock_router)
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
@@ -0,0 +1,4 @@
pytest>=7.0.0
pytest-cov>=4.0.0
pytest-asyncio>=0.21.0
aio-pika>=9.0.0
+210
View File
@@ -0,0 +1,210 @@
import unittest
from unittest.mock import AsyncMock, MagicMock, patch
from aio_pika.abc import (
AbstractIncomingMessage,
AbstractRobustChannel,
AbstractRobustQueue,
)
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQRoute, ServiceMessage
from amqp.model.service_message_type import ServiceMessageType
from amqp.router.router_base import (
CM_C_REQ_QUEUE,
CM_REQUEST_EXCHANGE,
)
from amqp.router.router_consumer import RouterConsumer
class TestRouterConsumer(unittest.IsolatedAsyncioTestCase):
def setUp(self):
# Create a mock for amq_adapter
self.amq_adapter = MagicMock()
self.amq_adapter.adapter_prefix.return_value = "test_"
self.amq_adapter.service_name = "test_service"
# Create the config mock with amq_adapter
self.config = MagicMock(spec=AMQConfiguration)
self.config.amq_adapter = self.amq_adapter
self.consumer = RouterConsumer(self.config)
# Mock channel
self.channel = AsyncMock(spec=AbstractRobustChannel)
self.consumer.channel = self.channel
# Mock queue
self.queue = AsyncMock(spec=AbstractRobustQueue)
self.channel.declare_queue.return_value = self.queue
def test_init(self):
"""Test initialization of RouterConsumer"""
self.assertIsNotNone(self.consumer)
self.assertEqual(self.consumer.amq_configuration, self.config)
@patch("amqp.router.router_consumer.logging_info")
async def test_init_routing_paths_consumer_success(self, mock_logging):
"""Test successful initialization of routing paths consumer"""
self.consumer.is_open = MagicMock(return_value=True)
await self.consumer.init_routing_paths_consumer(self.channel)
# Verify queue declaration
self.channel.declare_queue.assert_called_once_with(
name="test_" + CM_C_REQ_QUEUE, durable=True, exclusive=False, auto_delete=False
)
# Verify queue binding
self.queue.bind.assert_called_once_with(exchange=CM_REQUEST_EXCHANGE, routing_key="")
# Verify consumer setup
self.queue.consume.assert_called_once()
@patch("amqp.router.router_consumer.logging_error")
async def test_init_routing_paths_consumer_channel_closed(self, mock_logging):
"""Test initialization when channel is closed"""
self.consumer.is_open = MagicMock(return_value=False)
await self.consumer.init_routing_paths_consumer(self.channel)
mock_logging.assert_called_once()
self.channel.declare_queue.assert_not_called()
@patch("amqp.router.router_consumer.logging_debug")
@patch("amqp.router.router_consumer.logging_info")
async def test_handle_routing_data_request_message_valid(
self, mock_logging_info, mock_logging_debug
):
"""Test handling of valid routing data request message"""
# Setup
message = AsyncMock(spec=AbstractIncomingMessage)
message.delivery_tag = 1
message.consumer_tag = "test_consumer"
message.properties = MagicMock()
message.body = b"test_body"
# Mock service message factory
mock_service_message = ServiceMessage(
id=None,
payload_type=0,
message_type=ServiceMessageType.ROUTING_DATA_REQUEST,
recipient_name="test_service",
reply_to="other_service",
trace_info={},
message="test_payload",
)
self.consumer.service_message_factory = MagicMock()
self.consumer.service_message_factory.from_bytes = MagicMock(
return_value=mock_service_message
)
self.consumer.service_message_factory.of = MagicMock()
self.consumer.wrap_own_routes = MagicMock(return_value="test_routes")
self.consumer.publish_service_message = AsyncMock(return_value=True)
# Execute
await self.consumer.handle_routing_data_request_message(message)
# Verify
message.ack.assert_called_once_with(False)
self.consumer.publish_service_message.assert_called_once()
@patch("amqp.router.router_consumer.logging_info")
async def test_register_route_success(self, mock_logging):
"""Test successful route registration"""
# Setup
route = AMQRoute(
component_name="test_service",
exchange="test_exchange",
queue="test_queue",
key="test_key",
timeout=1000,
valid_until=0,
)
self.consumer.is_open = MagicMock(return_value=True)
self.consumer.publish_service_message = AsyncMock(return_value=True)
self.consumer.add_own_route = MagicMock()
# Execute
await self.consumer.register_route(route)
# Verify
self.consumer.publish_service_message.assert_called_once()
self.consumer.add_own_route.assert_called_once_with(route)
@patch("amqp.router.router_consumer.logging_warning")
async def test_register_route_channel_closed(self, mock_logging):
"""Test route registration when channel is closed"""
# Setup
route = AMQRoute(
component_name="test_service",
exchange="test_exchange",
queue="test_queue",
key="test_key",
timeout=1000,
valid_until=0,
)
self.consumer.is_open = MagicMock(return_value=False)
# Execute
await self.consumer.register_route(route)
# Verify
mock_logging.assert_called_once()
@patch("amqp.router.router_consumer.logging_info")
@patch("amqp.router.router_consumer.logging_warning")
async def test_remove_route_on_cleanup(self, mock_logging_warning, mock_logging_info):
"""Test route removal during cleanup"""
# Setup
route = AMQRoute(
component_name="test_service",
exchange="test_exchange",
queue="test_queue",
key="test_key",
timeout=1000,
valid_until=0,
)
self.consumer.publish_service_message = AsyncMock(return_value=True)
self.consumer.remove_route = MagicMock(return_value=True)
# Execute
result = await self.consumer.remove_route_on_cleanup(route)
# Verify
self.consumer.publish_service_message.assert_called_once()
self.consumer.remove_route.assert_called_once_with(route)
self.assertTrue(result)
@patch("amqp.router.router_consumer.logging_info")
@patch("amqp.router.router_consumer.logging_warning")
def test_cleanup(self, mock_logging_warning, mock_logging_info):
"""Test cleanup process"""
# Setup
routes = {
AMQRoute(
component_name="test_service",
exchange="test_exchange1",
queue="test_queue1",
key="test_key1",
timeout=1000,
valid_until=0,
),
AMQRoute(
component_name="test_service",
exchange="test_exchange2",
queue="test_queue2",
key="test_key2",
timeout=1000,
valid_until=0,
),
}
self.consumer.get_routes = MagicMock(return_value=routes)
self.consumer.remove_route_on_cleanup = AsyncMock(return_value=True)
self.consumer.wrap_routes = MagicMock(return_value="test_routes")
# Execute
self.consumer.cleanup()
# Verify
self.assertEqual(self.consumer.remove_route_on_cleanup.call_count, 2)
mock_logging_info.assert_called()
+164
View File
@@ -0,0 +1,164 @@
import unittest
from unittest.mock import AsyncMock, MagicMock, patch
from amqp.model.service_message_type import ServiceMessageType
from amqp.router.router_producer import RouterProducer
class DummyAMQAdapter:
def __init__(self):
self.service_name = "test-service"
self.route_mapping = "route-map"
def adapter_prefix(self):
return "prefix-"
class DummyConfig:
def __init__(self):
self.amq_adapter = DummyAMQAdapter()
class TestRouterProducer(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.config = DummyConfig()
self.router = RouterProducer(self.config)
def test_get_reply_to_queue_name(self):
self.assertEqual(self.router.get_reply_to_queue_name(), "prefix-P-REPLY-TO")
def test_get_response_queue_name(self):
self.assertEqual(self.router.get_response_queue_name(), "prefix-P-RESP")
def test_get_consume_queue_name(self):
self.assertEqual(self.router.get_consume_queue_name(), "prefix-C-AMQ")
def test_is_open_true(self):
channel = MagicMock()
channel.is_closed = False
self.assertTrue(self.router.is_open(channel))
def test_is_open_false(self):
channel = MagicMock()
channel.is_closed = True
self.assertFalse(self.router.is_open(channel))
self.assertFalse(self.router.is_open(None))
@patch("amqp.router.router_producer.logging_info")
@patch("amqp.router.router_producer.logging_error")
async def test_init_internal_routing_paths_channel_open(self, mock_error, mock_info):
channel = AsyncMock()
channel.is_closed = False
channel.declare_exchange = AsyncMock()
channel.declare_queue = AsyncMock()
queue = AsyncMock()
channel.declare_queue.return_value = queue
queue.bind = AsyncMock()
queue.consume = AsyncMock(return_value="ctag")
await self.router.init_internal_routing_paths(channel)
self.assertIsNotNone(self.router.cm_request_exchange)
self.assertIsNotNone(self.router.cm_response_exchange)
queue.bind.assert_awaited()
queue.consume.assert_awaited()
mock_error.assert_not_called()
@patch("amqp.router.router_producer.logging_error")
@patch("amqp.router.router_producer.logging_info")
async def test_init_internal_routing_paths_channel_closed(self, mock_info, mock_error):
channel = MagicMock()
channel.is_closed = True
await self.router.init_internal_routing_paths(channel)
mock_error.assert_called()
@patch("amqp.router.router_producer.logging_info")
@patch("amqp.router.router_producer.logging_debug")
async def test_handle_routing_data_message_valid(self, mock_debug, mock_info):
message = AsyncMock()
message.delivery_tag = 1
message.body = b"body"
message.message_id = "mid"
message.ack = AsyncMock()
# Patch service_message_factory
service_message = MagicMock()
service_message.payload_type = 0
service_message.is_valid.return_value = True
service_message.message_type = ServiceMessageType.ROUTING_DATA
service_message.reply_to = "other-service"
service_message.message = "route-string"
self.router.service_message_factory.from_bytes = MagicMock(return_value=service_message)
self.router.amq_configuration.amq_adapter.service_name = "test-service"
self.router.add_routes = MagicMock()
self.router.remove_routes = MagicMock()
await self.router.handle_routing_data_message(message)
message.ack.assert_awaited()
self.router.add_routes.assert_called_with("route-string")
self.router.remove_routes.assert_not_called()
@patch("amqp.router.router_producer.logging_info")
@patch("amqp.router.router_producer.logging_debug")
async def test_handle_routing_data_message_remove(self, mock_debug, mock_info):
message = AsyncMock()
message.delivery_tag = 1
message.body = b"body"
message.message_id = "mid"
message.ack = AsyncMock()
service_message = MagicMock()
service_message.is_valid.return_value = True
service_message.message_type = ServiceMessageType.ROUTING_DATA_REMOVE
service_message.payload_type = 0
service_message.reply_to = "other-service"
service_message.message = "route-string"
self.router.service_message_factory.from_bytes = MagicMock(return_value=service_message)
self.router.amq_configuration.amq_adapter.service_name = "test-service"
self.router.add_routes = MagicMock()
self.router.remove_routes = MagicMock()
await self.router.handle_routing_data_message(message)
message.ack.assert_awaited()
self.router.remove_routes.assert_called_with("route-string")
self.router.add_routes.assert_not_called()
@patch("amqp.router.router_producer.logging_info")
@patch("amqp.router.router_producer.logging_warning")
@patch("amqp.router.router_producer.logging_error")
async def test_request_routes_from_adapters_success(self, mock_error, mock_warning, mock_info):
self.router.cm_request_exchange = AsyncMock()
sm = MagicMock()
self.router.service_message_factory.of = MagicMock(return_value=sm)
self.router.publish_service_message = AsyncMock(return_value=True)
await self.router.request_routes_from_adapters()
mock_warning.assert_not_called()
mock_error.assert_not_called()
@patch("amqp.router.router_producer.logging_info")
@patch("amqp.router.router_producer.logging_warning")
@patch("amqp.router.router_producer.logging_error")
async def test_request_routes_from_adapters_fail(self, mock_error, mock_warning, mock_info):
self.router.cm_request_exchange = AsyncMock()
sm = MagicMock()
self.router.service_message_factory.of = MagicMock(return_value=sm)
self.router.publish_service_message = AsyncMock(return_value=False)
await self.router.request_routes_from_adapters()
mock_warning.assert_called()
mock_error.assert_not_called()
@patch("amqp.router.router_producer.logging_info")
async def test_publish_service_message_open(self, mock_info):
self.router.channel = MagicMock()
self.router.channel.is_closed = False
self.router.service_message_factory.serialize = MagicMock(return_value=b"bin")
exchange = AsyncMock()
exchange.publish = AsyncMock()
message = MagicMock()
with patch("amqp.router.router_producer.Message") as MockMessage:
MockMessage.return_value = MagicMock()
result = await self.router.publish_service_message(message, exchange)
exchange.publish.assert_awaited()
self.assertTrue(result)
async def test_publish_service_message_closed(self):
self.router.channel = MagicMock()
self.router.channel.is_closed = True
exchange = AsyncMock()
message = MagicMock()
result = await self.router.publish_service_message(message, exchange)
self.assertFalse(result)
+184
View File
@@ -0,0 +1,184 @@
import unittest
from unittest.mock import AsyncMock, Mock, patch
from amqp.config.amq_configuration import AMQAdapter, AMQConfiguration
from amqp.service.amq_service import AMQService
class TestAMQService(unittest.IsolatedAsyncioTestCase):
def setUp(self):
# Create a mock configuration
self.mock_adapter = Mock(spec=AMQAdapter)
self.mock_adapter.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
@@ -0,0 +1,12 @@
class AMQConfigurationTest:
def test_amqadapter(self):
assert False
def test_dispatch(self):
assert False
def test_backpressure(self):
assert False
def test_amqconfiguration(self):
assert False