Compare commits

...

84 Commits

Author SHA1 Message Date
Stanislav Hejny b41ca04801 feat/#61 - add one way communication option
Unit test coverage / pytest (push) Successful in 1m23s
/ build-and-push (push) Successful in 1m31s
Unit test coverage / pytest (pull_request) Successful in 1m8s
2025-07-23 18:33:34 +01:00
Stanislav Hejny 3ee62b0437 feat/#58 - fix unit test and deployment pipeline
Unit test coverage / pytest (push) Successful in 1m20s
/ build-and-push (push) Successful in 1m24s
2025-07-21 10:49:28 +01:00
Stanislav Hejny cf1804d23f feat: implement send_rpc method in DataMessageHandler class
Co-authored-by: aider (openrouter/anthropic/claude-3.7-sonnet) <aider@aider.chat>
2025-07-18 10:29:08 +01:00
Stanislav Hejny 04f1483c06 Merge branch 'feat-58-backpressure-reactive-streams' into feat/61-multi-user-support
Unit test coverage / pytest (push) Successful in 1m22s
/ build-and-push (push) Successful in 1m25s
2025-07-18 09:53:07 +01:00
Stanislav Hejny 050b5ddba4 fix: ensure AsyncMock returns awaitable channel in test_initialize
Unit test coverage / pytest (push) Successful in 1m37s
Unit test coverage / pytest (pull_request) Successful in 1m46s
/ build-and-push (push) Successful in 1m46s
Build and Publish Docker Image / build-and-push (push) Has been cancelled
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-18 09:29:54 +01:00
Stanislav Hejny a6ca319e92 fix: use AsyncMock for mock_channel in test setup
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-18 09:25:34 +01:00
Stanislav Hejny f2cc138540 feat/#50 - remove obsolete code
Build and Publish Docker Image / build-and-push (push) Waiting to run
Unit test coverage / pytest (push) Failing after 1m24s
Unit test coverage / pytest (pull_request) Failing after 1m56s
/ build-and-push (push) Successful in 2m0s
2025-07-18 07:59:20 +01:00
Stanislav Hejny f02823e18d feat/#61 - add client code for User Managemengt Service
Unit test coverage / pytest (push) Failing after 1m11s
/ build-and-push (push) Successful in 1m28s
2025-07-18 07:54:55 +01:00
Stanislav Hejny 5ab9406430 feat: implement UserManagementServiceClient with RabbitMQ communication
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-18 07:47:21 +01:00
Stanislav Hejny a6c918fe43 feat/58 - simplify complex logic
Build and Publish Docker Image / build-and-push (push) Waiting to run
Unit test coverage / pytest (push) Failing after 1m32s
Unit test coverage / pytest (pull_request) Failing after 1m38s
/ build-and-push (push) Successful in 1m46s
2025-07-17 17:51:01 +01:00
Stanislav Hejny 9cae4b168e fix: properly mock nested structure in AMQConfiguration for tests
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-17 10:49:47 +01:00
Stanislav Hejny 281004798e test: add comprehensive unit tests for DataMessageHandler
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-17 10:45:16 +01:00
Stanislav Hejny ef5a0363cd test: fix async mock in backpressure tests for coroutine compatibility
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-17 10:24:43 +01:00
Stanislav Hejny 7f9c02e458 test: simplify and enhance AMQService backpressure tests 2025-07-17 10:24:42 +01:00
Stanislav Hejny cc4fbc68fd test: add unit tests for BackpressureHandler and AMQService classes
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-16 18:48:18 +01:00
Stanislav Hejny 9c0cb3a334 refactor: update backpressure logic to handle availability correctly
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-16 18:04:09 +01:00
Stanislav Hejny f8fd9f24f8 feature/58 - update unit tests for coverage
Unit test coverage / pytest (push) Successful in 1m39s
Unit test coverage / pytest (pull_request) Successful in 1m46s
/ build-and-push (push) Successful in 1m52s
Build and Publish Docker Image / build-and-push (push) Has been cancelled
2025-07-15 17:12:26 +01:00
Stanislav Hejny c63941e4a1 fix: prevent unnecessary retries on last attempt in ID generation
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-15 16:55:07 +01:00
Stanislav Hejny 7cd2960fda feat/58 - improve Backpressure algorithm 2025-07-15 16:40:02 +01:00
Stanislav Hejny 7d4e358ce5 test: add AsyncMock to test backpressure handler conditions
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-12 23:02:19 +01:00
Stanislav Hejny 0e294b223d Fix the global unique ID generator recursive constructor call 2025-07-12 22:59:21 +01:00
Stanislav Hejny dd590a11b5 fix: correct consul_initial_retry_delay configuration logic
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-12 10:45:46 +01:00
Stanislav Hejny 5d8f07be3c refactor: prioritize env vars over config file in AMQAdapter settings
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-12 10:45:10 +01:00
Stanislav Hejny bd4272c30e feat: add Consul configuration to AMQAdapter class 2025-07-12 10:45:09 +01:00
Stanislav Hejny dbfd5ea50d test: add unit tests for backpressure handler methods and conditions
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-11 23:58:58 +01:00
Stanislav Hejny 826c859b32 feat: add IDLE event detection to check_overload_condition function
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-11 23:54:57 +01:00
Stanislav Hejny 8bfe6a4a49 feat: add update_backpressure_value method to BackpressureHandler class
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-11 23:48:31 +01:00
Stanislav Hejny 3150dc9907 feat: add toggle for CPU monitoring in backpressure handler
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-11 23:38:01 +01: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
Stanislav Hejny c387a2494e #23 refactor for cleverBRAG, fix failing unit tests
Unit test coverage / pytest (push) Failing after 1m10s
Unit test coverage / pytest (pull_request) Failing after 1m9s
2025-05-12 00:46:18 +01:00
Stanislav Hejny 27a930b3c3 #23 refactor to use shared code between services, part 2
Unit test coverage / pytest (push) Failing after 1m20s
Unit test coverage / pytest (pull_request) Failing after 1m18s
2025-05-11 23:57:27 +01:00
Stanislav Hejny 2f6dbd2ad8 #23 refactor to use shared code between services
Unit test coverage / pytest (pull_request) Failing after 1m9s
Unit test coverage / pytest (push) Failing after 1m11s
2025-05-11 15:06:39 +01:00
Stanislav Hejny 81e9616072 Merge branch 'develop' into cleverbrag-integration 2025-05-09 22:55:41 +01:00
stanislav.hejny 197671038b Merge pull request '#43 - fix the issues reported by flake8' (#44) from code-formatting into develop
Unit test coverage / pytest (push) Failing after 1m5s
CI for pypl publish / publish-lib (push) Failing after 1m3s
Reviewed-on: #44
2025-05-08 23:02:35 +00:00
Stanislav Hejny 15da316a1f #43 - fix the issues reported by flake8
Unit test coverage / pytest (push) Failing after 56s
Unit test coverage / pytest (pull_request) Failing after 1m1s
2025-05-08 23:52:33 +01:00
stanislav.hejny a2bb270535 Merge pull request '#40 - ensure the JSON format is handled correctly at all levels' (#42) from errors-as-json into develop
Unit test coverage / pytest (push) Failing after 1m19s
CI for pypl publish / publish-lib (push) Successful in 1m26s
Reviewed-on: #42
2025-05-07 12:07:32 +00:00
Stanislav Hejny 5bb6597b57 #40 - ensure the JSON format is handled correctly at all levels
Unit test coverage / pytest (push) Failing after 1m1s
Unit test coverage / pytest (pull_request) Failing after 1m18s
2025-05-07 13:05:04 +01:00
stanislav.hejny 3316d0891e Merge pull request 'errors-as-json' (#41) from errors-as-json into develop
Unit test coverage / pytest (push) Failing after 1m16s
CI for pypl publish / publish-lib (push) Successful in 1m23s
Reviewed-on: #41
2025-05-05 15:12:44 +00:00
Stanislav Hejny c862b3cc77 #40 - bump up wheel file version
Unit test coverage / pytest (push) Failing after 57s
Unit test coverage / pytest (pull_request) Failing after 1m20s
2025-05-05 16:11:15 +01:00
Stanislav Hejny d7b21af028 #40 - ensure return is a JSON with 'detail' field in case of the error 2025-05-05 15:55:47 +01:00
stanislav.hejny 687b134326 Merge pull request '#12 - add comments to properties to complete the documentation of AMQ Adapter' (#39) from config-documemntation into develop
Unit test coverage / pytest (push) Failing after 1m15s
CI for pypl publish / publish-lib (push) Failing after 1m16s
Reviewed-on: #39
2025-05-05 12:25:47 +00:00
Stanislav Hejny 5fbbdb9f83 #23 - first draft of CleverBRAG facade with endpoints discovery, suggested type (de)serialization and sample invocation 2025-04-28 21:59:22 +01:00
85 changed files with 8851 additions and 1414 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
+4
View File
@@ -0,0 +1,4 @@
[flake8]
max-line-length = 88
ignore = E501, W503, E203
exclude = .git,__pycache__,docs/source/conf.py,old,build,dist,.venv
+73
View File
@@ -0,0 +1,73 @@
name: Build and Publish Docker Image
#on:
# push:
# branches: [ main, master ]
# tags: [ 'v*' ]
# pull_request:
# branches: [ main, master ]
on:
push:
branches:
- master # publish release on master
- develop # publish snapshot on develop
- feat-58-backpressure-reactive-streams
workflow_dispatch:
# allow manual trigger
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
steps:
- name: Checkout code
uses: actions/checkout@v4
with:
ref: ${{ github.ref_name }}
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v2
- name: Login to Docker Registry
uses: docker/login-action@v3
with:
registry: ${{ env.REGISTRY_URL }}
username: ${{ secrets.REGISTRY_USER }}
password: ${{ secrets.REGISTRY_PASSWORD }}
if: github.event_name != 'pull_request'
- name: Extract metadata for Docker
id: meta
uses: docker/metadata-action@v4
with:
images: ${{ env.REGISTRY_URL }}/clevermicro/amq-adapter-python-demo
tags: |
type=ref,event=branch
type=ref,event=pr
type=semver,pattern={{version}}
type=semver,pattern={{major}}.{{minor}}
type=sha,format=short
- name: Build and push Docker image
uses: docker/build-push-action@v6
with:
context: .
file: ./Dockerfile
# push: ${{ github.event_name != 'pull_request' }}
push: true
tags: ${{ env.REGISTRY_URL }}/${{ env.REPOSITORY }}:latest
labels: ${{ steps.meta.outputs.labels }}
cache-from: type=gha
cache-to: type=gha,mode=max
+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 }}:20250709-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*
+21
View File
@@ -0,0 +1,21 @@
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.2.0
hooks:
- id: check-yaml
- id: end-of-file-fixer
- id: trailing-whitespace
- repo: https://github.com/PyCQA/isort
rev: 6.0.1
hooks:
- id: isort
args: ["--profile", "black"]
- repo: https://github.com/psf/black
rev: 25.1.0
hooks:
- id: black
language_version: python3.11
- repo: https://github.com/PyCQA/flake8
rev: 7.2.0
hooks:
- id: flake8
+49
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
```
+148 -163
View File
@@ -3,8 +3,8 @@ import json
import time
from asyncio import AbstractEventLoop
from threading import Thread
from typing import Optional
import psutil
from aio_pika import Message
from aio_pika.abc import AbstractRobustChannel
@@ -20,15 +20,15 @@ class ScaleRequestV1:
self,
serviceId: str,
taskId: str,
maxAvailability: int,
currentAvailability: int,
max_availability: int,
current_availability: int,
requestType: ScalingRequestAlert,
):
self.version = 1 # Version of the request, currently 1
self.serviceId = serviceId
self.taskId = taskId
self.maxAvailability = maxAvailability
self.currentAvailability = currentAvailability
self.max_availability = max_availability
self.current_availability = current_availability
self.requestType = requestType
self.thread = None
@@ -39,56 +39,23 @@ class ScaleRequestV1:
"version": self.version,
"serviceId": self.serviceId,
"taskId": self.taskId,
"maxAvailability": self.maxAvailability,
"currentAvailability": self.currentAvailability,
"maxAvailability": self.max_availability,
"currentAvailability": self.current_availability,
"requestType": self.requestType.value, # Using .value for Enum serialization
},
indent=2,
)
class RunningAverage:
def __init__(self, num_items):
self.buffer = [0] * num_items # Fixed-size array
self.pointer = 0 # Current position in array
self.num_items = num_items # Maximum number of items to store
self.is_filled = False # Track if buffer is fully filled
def add(self, value):
# Store the value at current pointer position
self.buffer[self.pointer] = value
# Move pointer to next position with wrap-around
self.pointer += 1
if self.pointer >= self.num_items:
self.pointer = 0
self.is_filled = True
def get_average(self):
# Determine how many items we should average
count = self.num_items if self.is_filled else self.pointer
if count == 0:
return 0.0 # Avoid division by zero
return sum(self.buffer) / count
def get_current_values(self):
"""Returns the values in chronological order (oldest first)"""
if not self.is_filled:
return self.buffer[: self.pointer]
return self.buffer[self.pointer :] + self.buffer[: self.pointer]
class BackpressureHandler:
# Track the number of messages currently being processed
current_parallel_executions = 0
current_availability = 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
# _callback_list is used in unit tests to record the invoked callbacks
_callback_list = None
def __init__(
self,
@@ -103,147 +70,160 @@ class BackpressureHandler:
self.exchange = None
self.swarm_service_id = self.config.amq_adapter.swarm_service_id
self.swarm_task_id = self.config.amq_adapter.swarm_task_id
self.parallel_workers = self.config.backpressure.threshold_threads
self.average_cpu_usage = None
self.do_loop = -1
self._resource_usage_changed = 0
self._resource_average_value = 0
self._last_resource_max_value = 100 # Default max value for CPU usage
self.max_availability = -1
self.current_availability = 1
self.max_load = 1
self.current_load = 0
def increase_parallel_executions(self):
"""Increase the number of parallel executions"""
def increase_current_load(self):
"""Increase the number of parallel executions, or current_availability load"""
logging_info(
"Backpressure: Increase parallel executions, current=%s",
self.current_parallel_executions,
"Backpressure: Increase current load (%s) by 1",
self.current_load,
)
self.current_parallel_executions += 1
self.current_load += 1
def decrease_parallel_executions(self):
"""Decrease the number of parallel executions"""
def decrease_current_load(self):
"""Decrease the number of parallel executions, or current load"""
logging_info(
"Backpressure: Decrease parallel executions, current=%s",
self.current_parallel_executions,
"Backpressure: Decrease current load(%s) by 1",
self.current_load,
)
if self.current_parallel_executions > 0:
self.current_parallel_executions -= 1
if self.current_load > 0:
self.current_load -= 1
def update_parallel_executions(self, count: int):
"""Update the number of parallel executions"""
self.current_parallel_executions = count
def update_current_availability(self, count: int):
"""Update the number of parallel executions, or current load"""
self.current_availability = count
def update_last_data_message_time(self):
logging_info("Backpressure: Update last data message time")
def update_last_backpressure_event_time(self):
# logging_info("Backpressure: Update last data message time")
"""Update the last data message time"""
self.last_data_message_time = time.time()
self.last_backpressure_event_time = time.time()
def start_backpressure_monitor(self) -> Thread:
async def update_backpressure_value(self, current_availability: int, maximum: int):
"""
Update the current backpressure value and check for overload conditions.
This method is called by the AMQService.backpressure() method to update
the current parallel executions count and check for overload conditions.
Args:
current_availability: Current value of the backpressure metric
maximum: Maximum value of the backpressure metric
"""
logging_info(
"Backpressure: Updating backpressure value, current_availability=%s, maximum=%s",
current_availability,
maximum,
)
if maximum > 0 and maximum > self.max_availability:
self.max_availability = maximum
# Update the current_availability / parallel executions count
self.update_current_availability(current_availability)
# Check for overload conditions
await self.check_overload_condition()
def start_backpressure_monitor(self) -> Optional[Thread]:
# Start the Backpressure monitor loop
self.thread = Thread(target=self.backpressure_monitor_loop)
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()
"""
Check if the current_availability availability is too low (OVERLOAD)
or if the service has been idle for too long with high availability (IDLE).
Note: current_availability represents available capacity, not used capacity.
- Low availability (close to 0) means OVERLOAD
- High availability (close to maximum) with no activity means IDLE
"""
current_time = time.time()
last_event_delta: float = current_time - self.last_backpressure_event_time
logging_info(
"Backpressure: Check overload condition, current=%s, max=%s",
self.current_parallel_executions,
self.parallel_workers,
"Backpressure: Check conditions, current_availability=%s, max=%s, last_activity=%s",
self.current_availability,
self.max_availability,
last_event_delta,
)
if self.current_parallel_executions >= self.parallel_workers - 1:
# Check for OVERLOAD condition - low availability (less than 20% of maximum)
if self.current_availability <= round(0.1 * self.max_availability):
# Check if the last backpressure event was not an overload
if self.last_backpressure_event != ScalingRequestAlert.OVERLOAD:
# Trigger the overload event
await self.handle_backpressure_overload_event()
self.last_backpressure_event_time = time.time()
self.last_backpressure_event = ScalingRequestAlert.OVERLOAD
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
_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
)
)
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
)
self.last_backpressure_event != ScalingRequestAlert.OVERLOAD
or last_event_delta > self.config.backpressure.idle_duration
):
# Trigger the overload event
await self.handle_backpressure_overload_event()
self.last_backpressure_event_time = _now
# update / reset time-window so that the OVERLOAD is not sent too often
self.last_backpressure_event_time = current_time
self.last_backpressure_event = ScalingRequestAlert.OVERLOAD
elif (
_cpu_usage_state == CPU_IDLE
and _now - _cpu_usage_changed > _idle_duration_millis
and _now - self.last_data_message_time > _idle_duration_millis
and (
_now - self.last_backpressure_event_time > _time_window_millis
or self.last_backpressure_event != ScalingRequestAlert.IDLE
)
# Check for IDLE condition - high availability (more than 80% of maximum) with no activity
elif self.current_availability >= round(0.8 * self.max_availability):
idle_duration = self.config.backpressure.idle_duration
# Check if service has been idle for longer than the configured duration
if (
last_event_delta > idle_duration
and self.last_backpressure_event == ScalingRequestAlert.IDLE
):
logging_info(
"Backpressure: Service has been idle for %s seconds (threshold: %s seconds)",
last_event_delta,
idle_duration,
)
# Trigger the idle event
await self._handle_backpressure_idle_event()
self.last_backpressure_event = ScalingRequestAlert.IDLE
self.last_backpressure_event_time = _now
# update / reset time-window so that the IDLE is not sent too often
self.last_backpressure_event_time = current_time
else:
# Trigger update event
if _now - self.last_backpressure_event_time > _time_window_millis:
# Trigger the update event
# Don't send IDLE right away when the availability increases, but wait for a while, send UPDATE instead
self.last_backpressure_event = ScalingRequestAlert.IDLE
_scaling_request: ScaleRequestV1 = ScaleRequestV1(
self.swarm_service_id,
self.swarm_task_id,
100,
_average_cpu_usage,
self.max_availability,
self.current_availability, # Current availability is passed directly
ScalingRequestAlert.UPDATE,
)
# Address the message to any adapter capable of supporting BACKPRESSURE request
await self.publish_backpressure_request(_scaling_request)
self.last_backpressure_event_time = _now
self.last_backpressure_event_time = current_time
# If neither OVERLOAD nor IDLE, and it's time for an update, send UPDATE event
elif last_event_delta > self.config.backpressure.time_window:
_scaling_request: ScaleRequestV1 = ScaleRequestV1(
self.swarm_service_id,
self.swarm_task_id,
self.max_availability,
self.current_availability, # Current availability is passed directly
ScalingRequestAlert.UPDATE,
)
# Address the message to any adapter capable of supporting BACKPRESSURE request
await self.publish_backpressure_request(_scaling_request)
self.last_backpressure_event_time = current_time
self.last_backpressure_event = ScalingRequestAlert.UPDATE
self.update_last_backpressure_event_time()
async def _backpressure_monitor(self):
"""Periodically monitor the backpressure conditions and trigger events accordingly"""
_monitor_interval = 0.5 # Monitor every 500ms
while self._do_loop():
await self.check_overload_condition()
await asyncio.sleep(_monitor_interval)
_loop.run_until_complete(_backpressure_monitor())
def backpressure_monitor_loop(self):
_loop = asyncio.new_event_loop()
_loop.run_until_complete(self._backpressure_monitor())
async def handle_backpressure_overload_event(self):
logging_warning("Backpressure: Capacity close to depleted!")
@@ -251,14 +231,12 @@ class BackpressureHandler:
scaling_request: ScaleRequestV1 = ScaleRequestV1(
self.swarm_service_id,
self.swarm_task_id,
self.parallel_workers,
self.parallel_workers - self.current_parallel_executions,
self.max_availability,
self.current_availability, # Current availability is passed directly
ScalingRequestAlert.OVERLOAD,
)
# Address the message to any management service instance capable of supporting BACKPRESSURE request
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.")
@@ -266,14 +244,12 @@ class BackpressureHandler:
scaling_request: ScaleRequestV1 = ScaleRequestV1(
self.swarm_service_id,
self.swarm_task_id,
self.parallel_workers,
self.parallel_workers - self.current_parallel_executions,
self.max_availability,
self.current_availability, # Current availability is passed directly
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
@@ -283,15 +259,13 @@ class BackpressureHandler:
if not self.exchange:
async def _wrap_rabbit_mq_api_init(channel):
_exchange = await channel.get_exchange(
name="cleverthis.clevermicro.management"
)
_exchange = await channel.get_exchange(name="cleverthis.clevermicro.management")
return _exchange
if BackpressureHandler._callback_list is not None:
BackpressureHandler._callback_list["_wrap_rabbit_mq_api_init"] = 1
self.exchange = await await_result(
asyncio.run_coroutine_threadsafe(
_wrap_rabbit_mq_api_init(self.channel), self.loop
)
asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api_init(self.channel), self.loop)
)
if self.exchange:
@@ -319,6 +293,17 @@ class BackpressureHandler:
return True
return False
await await_future(
asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api(), self.loop)
)
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))
def _do_loop(self) -> bool:
"""
Helper function for unit tests to perform several loops only. 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
+530 -170
View File
@@ -4,53 +4,71 @@ Methods here will invoke the actual CleverThis Service code, which are passed as
"""
import asyncio
import enum
import inspect
import json
import pathlib
import re
from asyncio import AbstractEventLoop
from typing import Dict, List, Callable, Coroutine, Any
from collections import defaultdict
from typing import Any, Callable, Coroutine, Dict, List
from aio_pika.abc import (
AbstractExchange,
AbstractRobustChannel,
AbstractRobustConnection,
AbstractExchange,
)
from aiohttp import ClientSession, ClientResponse
from dataclasses import dataclass
from fastapi import HTTPException
from aiohttp import ClientResponse, ClientSession
from fastapi import HTTPException, UploadFile
from fastapi.params import File
from fastapi.routing import APIRoute
from fastapi.security import OAuth2PasswordRequestForm
from pydantic import BaseModel
from starlette.datastructures import Headers
from starlette.responses import FileResponse
from amqp.adapter import serializer
from amqp.adapter.data_parser import AMQDataParser
from amqp.adapter.data_response_factory import DataResponseFactory
from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
from amqp.adapter.logging_utils import logging_error, logging_debug, 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 DataMessage, DataResponse
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()
@dataclass
class AMQMessage(DataMessage):
async def _create_reply_exchange_and_queue(
connection: AbstractRobustConnection,
) -> tuple[AbstractRobustChannel, AbstractExchange, str]:
"""
Expand the core DataMessage type to include the channel, which is required to create a dedicated queue
in case the response should initiate file download.
Create objects for reply.
Return (channel for reply queue, exchange for reply, reply queue name)
"""
def __init__(self, data_message: DataMessage, connection: AbstractRobustConnection):
# Copy all fields from DataMessage
self.magic = data_message.magic
self.id = data_message.id
self.method = data_message.method
self.domain = data_message.domain
self.path = data_message.path
self.headers = data_message.headers
self.trace_info = data_message.trace_info
self.base64_body = data_message.base64_body
# Add the new field
self.connection = connection
self.extra_data = (
data_message.extra_data if hasattr(data_message, "extra_data") else {}
# 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(
@@ -61,41 +79,22 @@ async def _convert_file_response(
) -> str:
"""
Converts a FileResponse object to a JSON string containing the filename and stream name.
The file content itself is transmitted through the named stream (which is implemented as temporary dedicated queue.
The file content itself is transmitted through the named stream (which is implemented as
temporary dedicated queue.
"""
async def _wrap_amq_api_calls(
connection: AbstractRobustConnection,
) -> tuple[AbstractRobustChannel, AbstractExchange, str]:
# Wrap calls to RabbitMQ API to make it Awaitable, so it can be run thread-safe with correct event loop
# This is needed because the RabbitMQ API is not thread-safe and needs to be called in the context of
# the event loop the first RabbitMQ connection was created with.
_channel = await connection.channel()
_exchange_name: str = AMQ_REPLY_TO_EXCHANGE
_exchange = await _channel.get_exchange(name=_exchange_name, ensure=True)
_queue = await _channel.declare_queue(
exclusive=False
) # Declare a unique, non-exclusive queue
_queue_name = _queue.name
await _queue.bind(
exchange=_exchange_name,
routing_key=_queue_name, # Use queue name as routing key
)
logging_debug(f"Publishing to exchange: {_exchange_name}, q: {_queue_name}")
return _channel, _exchange, _queue_name
# All calls to RabbitMQ API have to be called threadsafe, because those appear to be attached to parent
# event loop, and need to execute at that event loop, not the current one
# allow coroutine to execute. calling _future.result() will block the entire thread and the coroutine will never execute
# 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_availability one
# allow coroutine to execute. calling _future.result() will block the entire thread
# and the coroutine will never execute
_channel, _exchange, _routing_key = await await_result(
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(
_exchange, obj.path, _routing_key, loop
)
# as above, but luckily this time we don't need to wait for the result
_future = asyncio.run_coroutine_threadsafe(_channel.close(), loop)
asyncio.run_coroutine_threadsafe(_channel.close(), loop)
return json.dumps(
{
"files": [
@@ -112,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("]")
@@ -121,20 +122,28 @@ def is_likely_json(text: str) -> bool:
# ============= CleverThisServiceAdapter =============
class CleverThisServiceAdapter:
"""
A base class for CleverThis service adapter that handles the incoming AMQP messages and routes them to the
appropriate CleverThisService API endpoint.
IMPORTANT: This class is not intended to be used directly. It should be subclassed for each CleverThis service
and override the methods to handle the specific API calls.
A base class for CleverThis service adapter that handles the incoming AMQP messages and routes
them to the appropriate CleverThisService API endpoint.
IMPORTANT: This class is not intended to be used directly. It should be subclassed for
each CleverThis service and override the methods to handle the specific API calls.
"""
def __init__(
self, amq_configuration: AMQConfiguration, session: ClientSession = None
self,
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.
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
@@ -143,6 +152,81 @@ class CleverThisServiceAdapter:
self.require_authenticated_user = (
self.amq_configuration.amq_adapter.require_authenticated_user
)
self.routes = self.group_routes_by_method(api_routes=routes)
global g_extra_init_data
g_extra_init_data = extra_init_data or {}
def get_endpoint(self, endpoint: Any) -> Any:
"""
Optionally, Service can override this method to alter the endpoint mapping.
This is useful for example when the CleverThis service has a JSON wrapper
around the endpoint, like CleverSwarm, which might be the preferred endpoint to call.
:param endpoint: The original endpoint to be called.
:return: The actual endpoint to be called.
"""
return endpoint
def clone_and_adapt_route(self, original_route: APIRoute) -> APIRoute:
"""
Creates a clone of APIRoute. Checks if there's preferred endpoint with json serialization.
It also adjusts the path regex to ensure it matches the path in different API versions.
"""
# Need to adapt also the path regex to ensure it matches the path in different API versions
# Ensures the leading '^' is followed by any match and there's no '/' before final '$'
_new_regex_string = "^.*" + original_route.path_regex.pattern.lstrip("^").rstrip("/$") + "$"
_new_route: APIRoute = APIRoute(
path=original_route.path,
endpoint=self.get_endpoint(
original_route.endpoint
), # Service can alter the endpoint mapping here
response_model=original_route.response_model,
status_code=original_route.status_code,
tags=original_route.tags.copy() if original_route.tags else None,
dependencies=(
original_route.dependencies.copy() if original_route.dependencies else None
),
summary=original_route.summary,
description=original_route.description,
response_description=original_route.response_description,
responses=(original_route.responses.copy() if original_route.responses else None),
deprecated=original_route.deprecated,
name=original_route.name,
methods=original_route.methods.copy() if original_route.methods else None,
operation_id=original_route.operation_id,
response_model_include=original_route.response_model_include,
response_model_exclude=original_route.response_model_exclude,
response_model_by_alias=original_route.response_model_by_alias,
response_model_exclude_unset=original_route.response_model_exclude_unset,
response_model_exclude_defaults=original_route.response_model_exclude_defaults,
response_model_exclude_none=original_route.response_model_exclude_none,
include_in_schema=original_route.include_in_schema,
response_class=original_route.response_class,
dependency_overrides_provider=original_route.dependency_overrides_provider,
callbacks=(original_route.callbacks.copy() if original_route.callbacks else None),
openapi_extra=(
original_route.openapi_extra.copy() if original_route.openapi_extra else None
),
generate_unique_id_function=original_route.generate_unique_id_function,
)
_new_route.path_regex = re.compile(_new_regex_string)
return _new_route
def group_routes_by_method(self, api_routes: List[APIRoute]) -> Dict[str, List[APIRoute]]:
"""
Groups API routes by their HTTP method. This is important to ensure that URLPath is matched
correctly, as different methods may share the same URLPath (e.g. GET, POST PUT, DELETE).
"""
_method_to_routes = defaultdict(list) # Automatically initializes new lists for new keys
for _route in api_routes:
if (
isinstance(_route, APIRoute) and _route.methods
): # Ensure methods exist (should always be true for valid routes)
for method in _route.methods:
_method_to_routes[method].append(self.clone_and_adapt_route(_route))
return dict(_method_to_routes)
def get_content_type(self, message_headers: Dict[str, List[str]]) -> str:
"""
@@ -157,9 +241,11 @@ class CleverThisServiceAdapter:
async def get_current_user(self, token: str) -> object:
"""
To Be Overriden in subclass for given CleverThis Service, to return the current user based on the token,
in format appropriate for the service.
To Be Overriden in subclass for given CleverThis Service, to return the current_availability user based
on the token, in format appropriate for the service.
:param token: The token to be used for authentication.
:return: A user object or None if the user is not authenticated.
"""
return None
@@ -173,11 +259,9 @@ class CleverThisServiceAdapter:
_auth_user: Any | None = None
_amq_response: Any | None = None
logging_debug(
f"on_message: require_authenticated_user={self.require_authenticated_user}"
)
logging_debug(f"on_message: require_authenticated_user={self.require_authenticated_user}")
if self.require_authenticated_user:
_auth: str = message.headers.get("Authorization", "")
_auth: str = message.headers().get("Authorization", "")
logging_debug(f"Auth: {_auth}")
if isinstance(_auth, List):
_auth = _auth[0]
@@ -189,12 +273,8 @@ class CleverThisServiceAdapter:
except Exception as error:
logging_error(f"on_message: Error getting user: {error}")
if (
_auth_user
or not self.require_authenticated_user
or message.path.endswith("/login")
):
method = message.method.lower()
if _auth_user or not self.require_authenticated_user or message.path().endswith("/login"):
method = message.method().lower()
if method == "get":
_amq_response = await self.get(message, _auth_user)
elif method == "put":
@@ -206,160 +286,422 @@ class CleverThisServiceAdapter:
elif method == "delete":
_amq_response = await self.delete(message, _auth_user)
else:
raise ValueError(f"Unexpected HTTP method: {message.method}")
raise ValueError(f"Unexpected HTTP method: {message.method()}")
else:
_amq_response = DataResponseFactory.of(
message.id, 401, "text/plain", b"Unauthorized", message.trace_info
_amq_response = DataResponseFactory.of_error_message(
message.id(),
AMQErrorMessage(401, "Unauthorized", "Unauthorized"),
message.trace_info(),
)
logging_info(
f"ALL DONE in on_message id:{message.id}, resp code:{_amq_response.response_code}"
f"ALL DONE in on_message id:{message.id()}, resp code:{_amq_response.response_code()}"
)
return _amq_response
# ==================================================================================================
# ===================== C l e v e r T h i s A P I m e t h o d s ============================
# ==================================================================================================
async def get(self, message: DataMessage, auth_user: Any) -> DataResponse:
async def _authorize_forward_request(
self, message: DataMessage, request_coroutine, auth_user: Any
) -> DataResponse:
"""
Handles the GET requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
It will attempt to call the REST endpoint directly using HTTP protocol.
This method is used to authorize the request and forward it to the CleverThis Service API.
Used only for loose / 3rd party mode of coupling with the Service
"""
# AMQDataParser.parse_request_body(message)
# TODO: implement authorization logic here - ensure the security related headers are set
# This depends on what kind of security is required by the CleverThis service.
# However at this time there is no loosely coupled Service.
return await request_coroutine
# ==============================================================================================
# ===================== C l e v e r T h i s A P I m e t h o d s ========================
# ==============================================================================================
async def get(self, message: AMQMessage, auth_user: Any) -> DataResponse:
"""
Handles the GET requests for the CleverThis API. In tight coupling (self.session is None),
use path, path params and query string to determine the endpoint method to invoke.
In loose coupling (self.session is not None), forward the request as REST request using
the session, i.e. attempt to call the REST endpoint directly using HTTP protocol.
:param message: message from the AMQP that contains the GET request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
# unmatched path - try to invoke the service directly on this path
url = f"http://{self.service_host}:{self.service_port}{message.path}"
headers = {**message.headers, **message.trace_info}
if self.session is None:
return await self.handle_no_body_payload(message, auth_user, "GET")
# Forwards as REST - try to invoke the service directly on this path
url = f"http://{self.service_host}:{self.service_port}{message.path()}"
headers = {**message.headers(), **message.trace_info()}
logging_info(
f"SEND HEADERS out: {', '.join(f'{k}={v}' for k, v in headers.items())}, REQ={url}"
)
if self.session:
_http_resp: ClientResponse = await self.session.get(url, headers=headers)
return DataResponseFactory.of(
message.id,
message.id(),
_http_resp.status,
_http_resp.content_type,
await _http_resp.read(),
message.trace_info,
)
return DataResponseFactory.of(
message.id,
404,
"text/plain",
str("Destination location does not exists").encode("utf-8"),
message.trace_info,
message.trace_info(),
)
async def put(self, message: DataMessage, auth_user: Any) -> DataResponse:
async def put(self, message: AMQMessage, auth_user: Any) -> DataResponse:
"""
Handles the PUT requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
It will attempt to call the REST endpoint directly using HTTP protocol.
Handles the PUT requests for the CleverThis API. Uses path, path params, query string and
body payload to determine the endpoint method to invoke and its input arguments
In loose coupling (self.session is not None), attempt to call the REST endpoint directly
using HTTP protocol.
:param message: message from the AMQP that contains the PUT request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
return await self.handle_possible_form(
if self.session is None:
return await self.handle_body_payload(message, auth_user, "PUT")
# Forwards as REST - try to invoke the service directly on this path
return await self._authorize_forward_request(
message,
self.session.put(
f"http://{self.service_host}:{self.service_port}{message.path}",
data=message.body,
headers={**message.headers, **message.trace_info},
f"http://{self.service_host}:{self.service_port}{message.path()}",
data=message.body(),
headers={**message.headers(), **message.trace_info()},
),
auth_user,
)
async def post(self, message: DataMessage, auth_user: Any) -> DataResponse:
async def post(self, message: AMQMessage, auth_user: Any) -> DataResponse:
"""
Handles the POST requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
It will attempt to call the REST endpoint directly using HTTP protocol.
:param message: message from the AMQP that contains the POST request, including body and headers
Handles the POST requests for the CleverThis API. Uses path, path params, query string and
body payload to determine the endpoint method to invoke and its input arguments
In loose coupling (self.session is not None), attempt to call the REST endpoint directly
using HTTP protocol.
:param message: message from AMQP that contains the POST request, including body and headers
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
return await self.handle_possible_form(
if self.session is None:
return await self.handle_body_payload(message, auth_user, "POST")
return await self._authorize_forward_request(
message,
self.session.post(
f"http://{self.service_host}:{self.service_port}{message.path}",
data=message.body,
headers={**message.headers, **message.trace_info},
f"http://{self.service_host}:{self.service_port}{message.path()}",
data=message.body(),
headers={**message.headers(), **message.trace_info()},
),
auth_user,
)
async def handle_possible_form(
self, message: DataMessage, request_coroutine, auth_user: Any
async def head(self, message: AMQMessage, auth_user: Any) -> DataResponse:
"""
Handles the HEAD requests for the CleverThis API. In tight coupling (self.session is None),
use path, path params and query string to determine the endpoint method to invoke.
In loose coupling (self.session is not None), HEAD call is not supported.
:param message: message from AMQP containing the HEAD request, including body and headers
:param auth_user: user authenticated by the API Gateway, as UserSchema
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
if self.session is None:
return await self.handle_no_body_payload(message, auth_user, "HEAD")
raise NotImplementedError
async def delete(self, message: AMQMessage, auth_user: Any) -> DataResponse:
"""
Handles the DELETE requests for the CleverThis API.
:param message: message from AMQP that contains the DELETE request, including body & headers
:param auth_user: user authenticated by the API Gateway, as UserSchema
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
if self.session is None:
return await self.handle_no_body_payload(message, auth_user, "DELETE")
return await self._authorize_forward_request(
message,
self.session.post(
f"http://{self.service_host}:{self.service_port}{message.path()}",
data=message.body(),
headers={**message.headers(), **message.trace_info()},
),
auth_user,
)
async def options(self, message: AMQMessage, auth_user: Any) -> DataResponse:
"""
Handles the OPTIONS requests. In tight coupling (self.session is None),
use path, path params and query string to determine the endpoint method to invoke.
In loose coupling (self.session is not None), forward the request as REST request using
the session, i.e. attempt to call the REST endpoint directly using HTTP protocol.
:param message: message from AMQP that contains the OPTIONS request, including headers
:param auth_user: user authenticated by the API Gateway, as UserSchema
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
if self.session is None:
return await self.handle_no_body_payload(message, auth_user, "OPTIONS")
raise NotImplementedError
async def patch(self, message: AMQMessage, auth_user: Any) -> DataResponse:
"""
Handles the PATCH requests for the CleverBRAG API.
:param message: message from AMQP that contains the PATCH request, including body & headers
:param auth_user: user authenticated by the API Gateway, as UserSchema
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
if self.session is None:
return await self.handle_no_body_payload(message, auth_user, "PATCH")
return await self._authorize_forward_request(
message,
self.session.post(
f"http://{self.service_host}:{self.service_port}{message.path()}",
data=message.body(),
headers={**message.headers(), **message.trace_info()},
),
auth_user,
)
# ==============================================================================================
# ================= C l e v e r S w a r m p a y l o a d h a n d l e r s ==================
# ==============================================================================================
async def handle_no_body_payload(
self, message: AMQMessage, auth_user: Any, method: str
) -> DataResponse:
_args = AMQDataParser.parse_request_body(message)
return await request_coroutine
async def head(self, message: DataMessage, auth_user: Any) -> DataResponse:
"""
Handles the HEAD requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
It will attempt to call the REST endpoint directly using HTTP protocol.
:param message: message from the AMQP that contains the POST request, including body and headers
Handles the requests that do not have a body payload.
:param message: message from the AMQP that contains the request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
raise NotImplementedError
_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())
_signature_args = inspect.signature(_route.endpoint).parameters
if self.authenticated_user_arg_name in _signature_args.keys():
_args[self.authenticated_user_arg_name] = auth_user
# If the path matches, call the corresponding function
return await self.call_cleverthis_api(
message,
_args,
api_method=_route.endpoint,
success_code=_route.status_code,
)
async def delete(self, message: DataMessage, auth_user: Any) -> DataResponse:
return DataResponseFactory.of(
message.id(),
404,
"text/plain",
str(f"Destination location [{message.path()}] does not exists").encode("utf-8"),
message.trace_info(),
)
async def handle_body_payload(
self, message: AMQMessage, auth_user: Any, method: str
) -> DataResponse:
"""
Handles the DELETE requests for the CleverThis API. to be overriden in subclass. This is a base implementation.
It will attempt to call the REST endpoint directly using HTTP protocol.
:param message: message from the AMQP that contains the POST request, including body and headers
Handles the requests that have a body payload, including Multipart data.
:param message: message from the AMQP that contains the request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:param method: HTTP method (POST, PUT, etc.)
:return: DataResponse to be sent back to the API Gateway via AMQP
"""
raise NotImplementedError
_form_data = AMQDataParser.parse_request_body(message)
_routes = self.routes.get(method, [])
_path, _query_args = AMQDataParser.parse_get_url(message.path())
for _route in _routes:
_match = _route.path_regex.match(_path)
if _match:
# If the path matches, call the corresponding function
_form_data.update(_match.groupdict())
_form_data.update(_query_args)
for _body_param in _route.dependant.body_params:
if _body_param.name in _form_data:
_file_data = _form_data.get(_body_param.name, None)
actual_filepath = message.extra_data.get(_body_param.name, None)
_form_data[_body_param.name] = (
(
_body_param.type_.__name__ == "UploadFile"
or isinstance(_body_param.field_info, File)
)
and _file_data is not None
and actual_filepath is not None
and UploadFile(
file=pathlib.Path(actual_filepath).open("rb"),
size=_file_data.get("size"),
filename=_file_data.get("filename"),
headers=Headers({"Content-Type": _file_data.get("mediaType")}),
)
or _form_data[_body_param.name]
)
@staticmethod
async def call_cleverthis_api(
message: AMQMessage,
_verified_args = {}
_required_args = getattr(_route.endpoint, "__annotations__", {})
_signature_args = inspect.signature(_route.endpoint).parameters
for arg in _required_args:
if arg in _form_data:
_verified_args[arg] = _form_data[arg]
else:
# needed for CleverSwarm /login input
if _required_args[arg].__name__ == "OAuth2PasswordRequestForm":
try:
_verified_args[arg] = OAuth2PasswordRequestForm(**_form_data)
except Exception as e:
print(f"Error: {e}")
if self.authenticated_user_arg_name in _signature_args.keys():
_verified_args[self.authenticated_user_arg_name] = auth_user
return await self.call_cleverthis_api(
message,
_verified_args,
api_method=_route.endpoint,
success_code=_route.status_code if _route.status_code else 200,
)
return DataResponseFactory.of(
message.id(),
404,
"text/plain",
str(f"Destination location [{message.path()}] does not exists").encode("utf-8"),
message.trace_info(),
)
def _verify_and_instantiate_arguments(
self,
args: Any,
api_method: Callable[[Any, Any], Coroutine],
success_code: int,
loop: AbstractEventLoop | None = None,
) -> DataResponse:
) -> dict:
"""
Handles POST and PUT requests, which typically contain possibly complex payload data, passed as a dictionary.
Invokes the provided API method with the given arguments and authenticated user,
and returns the response as DataResponse.
:param message: DataMessage containing the details of the call and the payload
:param args: map of the arguments to pass to the API method
:param api_method: Callable - the method to invoke. Last parameter is the authenticated user
:param success_code: HTTP code to be returned when operation suceeds
:return: DataResponse class
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.
"""
try:
_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.annotation == int:
_verified_args[arg.name] = int(args[arg.name])
elif arg.annotation == float:
_verified_args[arg.name] = float(args[arg.name])
elif arg.annotation == bool:
_verified_args[arg.name] = bool(args[arg.name])
# TODO: more?
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:
_verified_args[arg.name] = args[arg.name]
_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:
logging_info(
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,
) -> DataResponse:
"""
Handles POST and PUT requests, which typically contain possibly complex payload data,
passed as a dictionary. Invokes the provided API method with the given arguments and
authenticated user, and returns the response as DataResponse.
:param message: DataMessage containing the details of the call and the payload
:param args: map of the arguments to pass to the API method
:param api_method: Callable - the method to invoke. Last parameter is the authenticated user
:param success_code: HTTP code to be returned when operation suceeds
:return: DataResponse class
"""
try:
_verified_args = self._verify_and_instantiate_arguments(
args=args, api_method=api_method
)
logging_info(
f"INVOKE ENDPOINT: {api_method.__name__}, ARGS: {', '.join(f'{k}={v}' for k, v in _verified_args.items())}"
"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 = (
@@ -367,12 +709,18 @@ class CleverThisServiceAdapter:
_result,
message.connection,
file_handler=StreamingFileHandler(),
loop=loop,
loop=self.loop,
),
"application/octet-stream",
)
elif isinstance(_result, BaseModel):
_json_result = _result.json()
_content_type = "application/json"
else:
if isinstance(_result, str):
_json_result = _result
else:
_json_result = python_type_to_json(_result)
_content_type = "application/json"
_binary_result = (
@@ -381,44 +729,56 @@ class CleverThisServiceAdapter:
else (
_json_result.body
if hasattr(_json_result, "body")
else (
_json_result.read()
if hasattr(_json_result, "read")
else _json_result
)
else (_json_result.read() if hasattr(_json_result, "read") else _json_result)
)
)
logging_info(
f"ALL DONE in call_cleverthis_api id:{message.id}, resp code:{success_code}"
f"ALL DONE in call_cleverthis_api id:{message.id()}, resp code:{success_code}"
)
return DataResponseFactory.of(
message.id,
message.id(),
success_code,
_content_type,
_binary_result,
message.trace_info,
message.trace_info(),
)
except HTTPException as http_error:
logging_error(f"Service REST endpoint invocation failed: {http_error}")
_content = str(http_error.detail)
if not is_likely_json(_content):
_content = '{"error": "' + _content + '"}'
logging_error(f"Service REST endpoint invocation failed: {http_error}")
return DataResponseFactory.of_error_message(
message.id(),
AMQErrorMessage(
http_error.status_code,
"Service Invocation Failed",
_content,
),
message.trace_info(),
)
return DataResponseFactory.of(
message.id,
message.id(),
http_error.status_code,
"application/json",
_content.encode("utf-8"),
message.trace_info,
message.trace_info(),
)
except Exception as error:
logging_error(f"Service endpoint invocation failed: {error}")
_content = str(error)
if not is_likely_json(_content):
_content = '{"error": "' + _content + '"}'
return DataResponseFactory.of_error_message(
message.id(),
AMQErrorMessage(
500,
"Service Invocation Failed",
_content,
),
message.trace_info(),
)
return DataResponseFactory.of(
message.id,
message.id(),
500,
"application/json",
_content.encode("utf-8"),
message.trace_info,
message.trace_info(),
)
@@ -0,0 +1,98 @@
import logging
import os
import random
import socket
import time
import consul_kv
from amqp.config.amq_configuration import AMQAdapter
logger = logging.getLogger(__name__)
def get_config_values(config: AMQAdapter) -> tuple:
"""
Get configuration values from AMQConfiguration or environment variables.
Returns:
tuple: (consul_host, consul_port, consul_counter_key, max_retries, retry_delay_seconds)
"""
try:
consul_host = config.consul_host
consul_port = config.consul_port
consul_counter_key = config.consul_counter_key
max_retries = config.consul_max_retries
retry_delay_seconds = config.consul_initial_retry_delay
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", consul_kv.settings.DEFAULT_ENDPOINT.split(":")[0]),
int(os.environ.get("CONSUL_PORT", consul_kv.settings.DEFAULT_ENDPOINT.split(":")[1])),
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.2)),
)
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(config: AMQAdapter) -> 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(config)
)
try:
c = consul_kv.Connection(endpoint=f"{consul_host}:{consul_port}", timeout=5)
index, data = c.get(consul_counter_key)
if data is None:
logger.info(f"Initializing Consul counter at {consul_counter_key}")
if c.put(consul_counter_key, "1", cas=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.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...")
if attempt < max_retries - 1:
time.sleep(retry_delay_seconds * (1 << attempt)) # Exponential backoff
index, data = c.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()
+51 -29
View File
@@ -5,21 +5,22 @@ methods to create, serialize and deserialize DataMessage objects.
import base64
import json
import re
from io import BytesIO
from datetime import datetime, timezone
from typing import Dict, List
from io import BytesIO
from typing import Any, Dict, List
import amqp.adapter.file_handler
from amqp.model.model import DataMessage
from amqp.model.snowflake_id import SnowflakeId
from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
from amqp.adapter.pydantic_serializer import python_type_to_json
from amqp.adapter.serializer import (
map_with_list_as_string,
map_as_string,
map_with_list_as_string,
parse_map_list_string,
parse_map_string,
)
from amqp.model.model import DataMessage, DataMessageMagicByte
from amqp.model.snowflake_id import SnowflakeId
from amqp.router.utils import sanitize_as_rabbitmq_name
from amqp.adapter.file_handler import StreamingFileHandler, FileHandler
class DataMessageFactory:
@@ -30,7 +31,6 @@ class DataMessageFactory:
MACHINE_ID_MASK = (1 << MACHINE_ID_BITS) - 1
SEQUENCE_MASK = (1 << SEQUENCE_BITS) - 1
MAGIC_BYTE_HTTP_REQUEST = "A"
_instance = None
snowflake_epoch = datetime(2024, 10, 1, 0, 0, tzinfo=timezone.utc).timestamp()
@@ -100,12 +100,10 @@ class DataMessageFactory:
:return: DataMessage object
"""
serializable_headers = headers.copy() if headers else {}
payload = base64.b64encode(
message.encode("utf-8") if isinstance(message, str) else message
)
payload = base64.b64encode(message.encode("utf-8") if isinstance(message, str) else message)
trace_info = {}
return DataMessage(
self.MAGIC_BYTE_HTTP_REQUEST,
DataMessageMagicByte.HTTP_REQUEST.value,
self.generate_snowflake_id(),
method,
domain,
@@ -115,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:
"""
@@ -122,7 +146,7 @@ class DataMessageFactory:
:param message: DataMessage input
:return: routing key (as string) compliant to RabbitMQ allowed character set.
"""
return sanitize_as_rabbitmq_name(f"{message.domain}.{message.path}")
return sanitize_as_rabbitmq_name(f"{message.domain()}.{message.path()}")
@staticmethod
def get_timestamp_from_id(_id: SnowflakeId) -> int:
@@ -132,9 +156,7 @@ class DataMessageFactory:
:return: the message timestamp
"""
bits = _id.get_bits()
low = bits[0] >> (
DataMessageFactory.SEQUENCE_BITS + DataMessageFactory.MACHINE_ID_BITS
)
low = bits[0] >> (DataMessageFactory.SEQUENCE_BITS + DataMessageFactory.MACHINE_ID_BITS)
high = bits[1] << (
64 - DataMessageFactory.SEQUENCE_BITS - DataMessageFactory.MACHINE_ID_BITS
)
@@ -149,13 +171,13 @@ class DataMessageFactory:
"""
return (
f"DataMessage{{"
f"id={message.id}, "
f"method='{message.method}', "
f"domain='{message.domain}', "
f"path='{message.path}', "
f"headers={{{map_with_list_as_string(message.headers)}}}, "
f"traceInfo={{{map_as_string(message.trace_info)}}}, "
f"base64body={message.base64_body.decode()}"
f"id={message.id()}, "
f"method='{message.method()}', "
f"domain='{message.domain()}', "
f"path='{message.path()}', "
f"headers={{{map_with_list_as_string(message.headers())}}}, "
f"traceInfo={{{map_as_string(message.trace_info())}}}, "
f"base64body={message.base64body().decode()}"
f"}}"
)
@@ -166,19 +188,19 @@ class DataMessageFactory:
:param message: message to serialize
:return: messages as bytes
"""
id_bytes = str(message.id).encode()
id_bytes = str(message.id()).encode()
prolog = (
f"|m={message.method}|d={message.domain}|p={message.path}|h={{{map_with_list_as_string(message.headers)}}}|t={{{map_as_string(message.trace_info)}}}|b="
f"|m={message.method()}|d={message.domain()}|p={message.path()}|h={{{map_with_list_as_string(message.headers())}}}|t={{{map_as_string(message.trace_info())}}}|b="
).encode("utf-8")
header_length = 1 + len(id_bytes) + len(prolog)
prolog_len = header_length.to_bytes(2, byteorder="big")
msg_length = 2 + header_length + len(message.base64_body)
msg_length = 2 + header_length + len(message.base64body())
with BytesIO(bytearray(msg_length)) as baos:
baos.write(prolog_len)
baos.write(DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST.encode("utf-8"))
baos.write(DataMessageMagicByte.HTTP_REQUEST.value.encode("utf-8"))
baos.write(id_bytes)
baos.write(prolog)
baos.write(message.base64_body)
baos.write(message.base64body())
return baos.getvalue()
@staticmethod
@@ -192,7 +214,7 @@ class DataMessageFactory:
metadata = input_bytes[2 : prolog_len + 2].decode()
pattern = re.compile(
f"{DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST}"
f"{DataMessageMagicByte.HTTP_REQUEST.value}"
r"([0-9A-Fa-f]+)\|"
r"m=([^|]*)\|"
r"d=([^|]*)\|"
@@ -220,7 +242,7 @@ class DataMessageFactory:
trace_info = parse_map_string(trace_info_str)
return DataMessage(
DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST,
DataMessageMagicByte.HTTP_REQUEST.value,
SnowflakeId.from_hex(id_str),
method,
domain,
+22 -29
View File
@@ -1,15 +1,15 @@
import io
import json
import python_multipart
from fastapi import UploadFile
from python_multipart.multipart import Field, File
from typing import Tuple, Dict
from typing import Dict, Tuple
from urllib.parse import parse_qs, urlparse
from xml.etree import ElementTree
from amqp.adapter.logging_utils import logging_error, logging_debug
from amqp.model.model import DataMessage
import python_multipart
from fastapi import UploadFile
from python_multipart.multipart import Field, File
from amqp.adapter.logging_utils import logging_debug, logging_error
from amqp.model.model import AMQMessage
class MultipartFormDataParser:
@@ -22,15 +22,14 @@ class MultipartFormDataParser:
included, the message is transmitted in the unchanged multipart/form-data format.
"""
def __init__(self, message: DataMessage):
def __init__(self, message: AMQMessage):
self.form_data: dict = {}
self.is_multipart = False
self.is_json = False
self.is_valid = True
# Convert each list of strings to a single string joined by newline, then to bytes
_headers = {
key: "\n".join(value).encode("utf-8")
for key, value in message.headers.items()
key: "\n".join(value).encode("utf-8") for key, value in message.headers().items()
}
try:
if (
@@ -41,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}")
@@ -52,15 +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
@@ -70,8 +68,9 @@ class MultipartFormDataParser:
def on_file(self, file: File):
logging_debug(str(file))
file.file_object.seek(0)
self.form_data[file.field_name.decode("utf-8")] = UploadFile(
file.file_object, size=file.size, filename=file.file_name
file.file_object, size=file.size, filename=file.file_name.decode("utf-8")
)
@@ -111,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
@@ -130,13 +128,11 @@ class AMQDataParser:
_parsed = urlparse(url)
_query_params = parse_qs(_parsed.query)
# Convert values from lists to single values when there's only one value
_simplified_params = {
k: v[0] if len(v) == 1 else v for k, v in _query_params.items()
}
_simplified_params = {k: v[0] if len(v) == 1 else v for k, v in _query_params.items()}
return _parsed.path, _simplified_params
@staticmethod
def parse_request_body(message: DataMessage):
def parse_request_body(message: AMQMessage):
"""
Parses the HTTP POST request body based on the MIME type.
This is typically called fom the concrete CleverThis Service to extract parameters from the message body.
@@ -146,14 +142,14 @@ class AMQDataParser:
Returns:
dict: Parsed data as a dictionary.
"""
_content_type = message.headers.get("Content-Type", "")
_content_type = message.headers().get("Content-Type", "")
if not _content_type:
raise ValueError("Content-Type header is required")
content_type = _content_type[0]
# Parse JSON
if content_type == "application/json":
try:
return json.loads(message.body())
return json.loads(message.body() or "{}")
except json.JSONDecodeError as e:
logging_error(f"Invalid JSON: {e}")
raise ValueError(f"Invalid JSON: {e}")
@@ -162,15 +158,12 @@ class AMQDataParser:
elif content_type == "application/x-www-form-urlencoded":
try:
_form_data: dict = {
k: v[0] if len(v) == 1 else v
for k, v in parse_qs(message.body()).items()
k: v[0] if len(v) == 1 else v for k, v in parse_qs(message.body()).items()
}
if len(_form_data) == 0 and len(message.body()) > 0:
_form_data = json.loads(message.body())
if _form_data["files"] and _form_data["form"]:
_form_data = MultipartFormDataParser.process_form_data(
_form_data
)
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:
logging_error(f"Invalid URL-encoded form data: {e}")
+44 -18
View File
@@ -5,13 +5,13 @@ method to create, serialize and deserialize DataResponse message
import base64
import re
from io import BytesIO
from amqp.model.model import DataResponse
from amqp.adapter.serializer import map_as_string, parse_map_string
from amqp.model.model import AMQErrorMessage, DataMessageMagicByte, DataResponse
from amqp.model.snowflake_id import SnowflakeId
from amqp.adapter.serializer import parse_map_string, map_as_string
class DataResponseFactory:
MAGIC_BYTE_HTTP_REQUEST = "A"
@staticmethod
def create_async_response_message(_id: SnowflakeId) -> DataResponse:
@@ -20,8 +20,14 @@ class DataResponseFactory:
:param _id: ID
:return: Response message
"""
return DataResponse(
str(_id), 200, "application/text", "OK".encode("utf-8"), "", "", {}
return DataResponseFactory.of(
id=_id,
response_code=200,
content_type="application/text",
body="OK".encode("utf-8"),
trace_info={},
error="",
error_cause="",
)
@staticmethod
@@ -31,6 +37,8 @@ class DataResponseFactory:
content_type: str,
body: bytes,
trace_info: dict[str, str],
error: str | None = None,
error_cause: str | None = None,
) -> DataResponse:
"""
Create the DataResponse message supplying all key values
@@ -42,13 +50,31 @@ class DataResponseFactory:
:return: DataResponseMessage object
"""
return DataResponse(
id=id.as_string(),
response_code=response_code,
magic=DataMessageMagicByte.HTTP_REQUEST.value,
id=id,
response_code=str(response_code),
content_type=content_type,
response=body,
error=None,
error_cause=None,
error=error,
error_cause=error_cause,
headers={},
trace_info=trace_info,
base64body=base64.b64encode(body),
)
@staticmethod
def of_error_message(
id: SnowflakeId,
error_message: AMQErrorMessage,
trace_info: dict[str, str],
) -> DataResponse:
return DataResponseFactory.of(
id=id,
response_code=error_message.response_code,
content_type=error_message.content_type,
body=error_message.serialize(),
trace_info=trace_info,
error=error_message.error,
error_cause=error_message.detail,
)
@staticmethod
@@ -58,19 +84,19 @@ class DataResponseFactory:
:param response: Object to serialize
:return: byte array as serialized message
"""
id_bytes = response.id.encode("utf-8")
id_bytes = response.id().as_string().encode("utf-8")
prolog = (
f"|r={response.response_code}|c={response.content_type}|e={response.error}|f={response.error_cause}|t={{{map_as_string(response.trace_info)}}}|b="
f"|r={response.response_code()}|c={response.content_type()}|e={response.error()}|f={response.error_cause()}|t={{{map_as_string(response.trace_info())}}}|b="
).encode("utf-8")
header_length = 1 + len(id_bytes) + len(prolog)
prolog_len = header_length.to_bytes(2, byteorder="big")
msg_length = 2 + header_length + len(response.response)
msg_length = 2 + header_length + len(response.base64body())
with BytesIO(bytearray(msg_length)) as baos:
baos.write(prolog_len)
baos.write(DataResponseFactory.MAGIC_BYTE_HTTP_REQUEST.encode("utf-8"))
baos.write(DataMessageMagicByte.HTTP_REQUEST.value.encode("utf-8"))
baos.write(id_bytes)
baos.write(prolog)
baos.write(base64.b64encode(response.response))
baos.write(response.base64body())
return baos.getvalue()
@staticmethod
@@ -84,7 +110,7 @@ class DataResponseFactory:
metadata = input_bytes[2 : prolog_len + 2].decode()
pattern = re.compile(
f"{DataResponseFactory.MAGIC_BYTE_HTTP_REQUEST}"
f"{DataMessageMagicByte.HTTP_REQUEST.value}"
r"([0-9A-Fa-f]+)\|"
r"r=([^|]*)\|"
r"c=([^|]*)\|"
@@ -110,14 +136,14 @@ class DataResponseFactory:
trace_info = parse_map_string(trace_info_str)
return DataResponse(
return DataResponseFactory.of(
id_str,
response_code,
content_type,
base64.b64decode(body_b64),
trace_info,
error,
error_cause,
trace_info,
)
+76 -70
View File
@@ -4,22 +4,22 @@ implements the file downloader for the amqp adapter, where file is sent in chunk
import asyncio
import json
import os
import logging
from asyncio import Future, AbstractEventLoop
import os
from asyncio import AbstractEventLoop
from collections import defaultdict
from pathlib import Path
from aio_pika import connect_robust, Message
from aio_pika import Message, connect_robust
from aio_pika.abc import (
AbstractIncomingMessage,
AbstractRobustChannel,
AbstractQueue,
AbstractExchange,
AbstractIncomingMessage,
AbstractQueue,
AbstractRobustChannel,
DeliveryMode,
)
from amqp.adapter.logging_utils import logging_error, logging_debug, logging_info
from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
from amqp.router.utils import await_result
# Global dictionary to track completed downloads
@@ -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):
@@ -166,7 +204,6 @@ class StreamingFileHandler(FileHandler):
message_data (dict): The entire message
output_dir (str): The directory where the file should be saved.
"""
global download_locations
filename = file_info["filename"]
size = file_info["size"]
stream_name = file_info["streamName"]
@@ -185,64 +222,31 @@ 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]:
"""
Asynchronously downloads a bytearay from RabbitMQ based on the provided file information.
It assumes entire array is in the first (and only) chunk
Args:
loop: The asyncio event loop.
rabbitmq_url: The RabbitMQ connection URL.
queue_name (str): The directory where the file should be saved.
Internal method for downloading a file from a queue to RAM.
Return (bytearray, eof status).
"""
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)
_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}"
)
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():
@@ -257,8 +261,21 @@ class StreamingFileHandler(FileHandler):
await _connection.close()
return _res, _eof
async def download_buffer(self, loop, rabbitmq_url, queue_name: str | None) -> (bytearray, int):
"""
Asynchronously downloads a bytearay from RabbitMQ based on the provided file information.
It assumes entire array is in the first (and only) chunk
Args:
loop: The asyncio event loop.
rabbitmq_url: The RabbitMQ connection URL.
queue_name (str): The directory where the file should be saved.
"""
logging_debug(f"Downloading buffer from stream: {queue_name}")
if queue_name:
_local_res, _local_eof = await await_result(
asyncio.run_coroutine_threadsafe(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
@@ -281,7 +298,6 @@ class StreamingFileHandler(FileHandler):
loop: The asyncio event loop of the RabbitMQ API
output_dir
"""
global download_locations
try:
_message_data = json.loads(message.body.decode())
_amq_message_data = json.loads(json_data.decode())
@@ -298,22 +314,16 @@ class StreamingFileHandler(FileHandler):
files_to_download.append(files_info)
if not files_to_download:
logging_info(
" [%s] No files to download in this message.", message.delivery_tag
)
logging_info(" [%s] No files to download in this message.", message.delivery_tag)
# await message.ack()
return defaultdict()
# Create a task for each file download
_tasks = []
for file_info in files_to_download:
logging_info(
f" [{message.delivery_tag}] about to create task for {file_info}"
)
logging_info(f" [{message.delivery_tag}] about to create task for {file_info}")
_task = asyncio.create_task(
self.download_file(
loop, rabbitmq_url, file_info, _message_data, output_dir
)
self.download_file(loop, rabbitmq_url, file_info, _message_data, output_dir)
)
_tasks.append(_task)
logging_info(
@@ -379,9 +389,7 @@ class StreamingFileHandler(FileHandler):
"total_size": _file_size,
"chunk_size": len(_chunk),
"eof": (
END_OF_FILE
if max_chunk_size + _offset >= _file_size
else IN_PROGRESS
END_OF_FILE if max_chunk_size + _offset >= _file_size else IN_PROGRESS
),
},
)
@@ -394,9 +402,7 @@ class StreamingFileHandler(FileHandler):
while not _future.done():
# sleep allows coroutine to execute. must wait for completion to ensure all chunks written in right order
await asyncio.sleep(0.02)
logging_debug(
f"Waiting for chunk {_offset}/{_file_size} future= {_future}"
)
logging_debug(f"Waiting for chunk {_offset}/{_file_size} future= {_future}")
_offset += len(_chunk)
logging_debug(
+25 -22
View File
@@ -4,23 +4,26 @@ import os
import sys
import threading
import traceback
from importlib.metadata import version
verbose = os.environ.get("AMQ_VERBOSE_LOGGING", "0") == "1"
# Get version from installed package metadata
__version__ = version("amq_adapter") # Name must match pyproject.toml [project] name
__verbose__ = os.environ.get("AMQ_VERBOSE_LOGGING", "0") == "1"
def get_context_info():
"""
Retrieves the current module, line number, and function name.
Retrieves the current_availability module, line number, and function name.
Returns:
A tuple containing (module name, line number, function name).
Returns (None, None, None) if it fails to retrieve the information.
"""
try:
# Get the frame object for the current function call
# Get the frame object for the current_availability function call
frame = inspect.currentframe()
if frame is not None:
# Get the frame object for the caller of the current function
# Get the frame object for the caller of the current_availability function
frame = frame.f_back.f_back
if frame:
code_context = frame.f_code
@@ -41,7 +44,7 @@ def get_context_info():
def logging_error(msg: str, *args) -> None:
"""
Log an error message according to current logging for 'ERROR' level.
Log an error message according to current_availability logging for 'ERROR' level.
Add module name, line and function name where the error occurred, on single line.
Args:
msg (str): The error message to log.
@@ -51,65 +54,65 @@ def logging_error(msg: str, *args) -> None:
if _exec_info:
_tb = traceback.extract_tb(sys.exc_info()[2])[-1]
logging.error(
f"{_thread_id} [E] '{msg}' ---> '{_tb.filename}':{_tb.lineno}, '{_tb.name}()'",
f"{_thread_id}:{__version__} [E] '{msg}' ---> '{_tb.filename}':{_tb.lineno}, '{_tb.name}()'",
*args,
)
else:
module, line_number, function_name = get_context_info()
if module and line_number and function_name:
logging.error(
f"{_thread_id} [E] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
f"{_thread_id}:{__version__} [E] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
*args,
)
else:
logging.error(f"{_thread_id} [E] {msg}", *args)
logging.error(f"{_thread_id}:{__version__} [E] {msg}", *args)
def logging_warning(msg: str, *args) -> None:
"""
Log a warning message according to current logging for 'WARN' level.
Log a warning message according to current_availability logging for 'WARN' level.
Add module name, line and function name where the print occurred, on single line.
Args:
msg (str): The warning message to log.
"""
_thread_id = threading.get_ident()
if not verbose:
logging.warning(f"{_thread_id} [W] {msg}", *args)
if not __verbose__:
logging.warning(f"{_thread_id}:{__version__} [W] {msg}", *args)
return
module, line_number, function_name = get_context_info()
if module and line_number and function_name:
logging.warning(
f"{_thread_id} [W] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
f"{_thread_id}:{__version__} [W] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
*args,
)
else:
logging.warning(f"{_thread_id} [W] {msg}", *args)
logging.warning(f"{_thread_id}:{__version__} [W] {msg}", *args)
def logging_info(msg: str, *args) -> None:
"""
Log an info message according to current logging for 'INFO' level.
Log an info message according to current_availability logging for 'INFO' level.
Add module name, line and function name where the print occurred, on single line.
Args:
msg (str): The info message to log.
"""
_thread_id = threading.get_ident()
if not verbose:
logging.info(f"{_thread_id} [*] {msg}", *args)
if not __verbose__:
logging.info(f"{_thread_id}:{__version__} [*] {msg}", *args)
return
module, line_number, function_name = get_context_info()
if module and line_number and function_name:
logging.info(
f"{_thread_id} [*] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
f"{_thread_id}:{__version__} [*] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
*args,
)
else:
logging.info(f"{_thread_id} [*] {msg}", *args)
logging.info(f"{_thread_id}:{__version__} [*] {msg}", *args)
def logging_debug(msg: str, *args) -> None:
"""
Log a debug message according to current logging for 'DEBUG' level.
Log a debug message according to current_availability logging for 'DEBUG' level.
Add module name, line and function name where the print occurred, on single line.
Args:
msg (str): The debug message to log.
@@ -117,9 +120,9 @@ def logging_debug(msg: str, *args) -> None:
_thread_id = threading.get_ident()
module, line_number, function_name = get_context_info()
if module and line_number and function_name:
logging.info(
f"{_thread_id} [DBG] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
logging.debug(
f"{_thread_id}:{__version__} [DBG] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
*args,
)
else:
logging.debug(f"{_thread_id} [DBG] {msg}", *args)
logging.debug(f"{_thread_id}:{__version__} [DBG] {msg}", *args)
+91 -81
View File
@@ -2,12 +2,15 @@
This module provides serialization and deserialization functions for Pydantic models that extend BaseModel
"""
import uuid
import importlib
import json
from datetime import datetime
from typing import Dict, Any
import uuid
from datetime import date, datetime
from types import UnionType
from typing import Any, Dict, Union, get_args, get_origin
from urllib.parse import parse_qs, urlparse
from pydantic import BaseModel
from urllib.parse import urlparse, parse_qs
from amqp.adapter.logging_utils import logging_debug
@@ -25,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):
@@ -46,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:
@@ -62,91 +111,52 @@ 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}")
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
# 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:
return v # Keep as string
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))
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
return cls(**obj_dict)
def parse_url_query(url):
+9 -2
View File
@@ -4,7 +4,7 @@ Helper functions used to aid with (de)serialization of messages sent over AMQP.
import struct
from io import BytesIO
from typing import List, Union, Dict
from typing import Dict, List, Union
def long_to_bytes(x: int) -> bytes:
@@ -18,7 +18,7 @@ def bytes_to_long(b: bytes) -> int:
def read_long(bis: BytesIO) -> int:
try:
return struct.unpack("q", bis.read(8))[0]
except 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:
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 :])
@@ -63,3 +66,7 @@ def parse_map_list_string(input_str: str) -> Dict[str, List[str]]:
for token in input_str[1:-1].split("],"):
add_to_map(map_, token)
return map_
def str_to_bool(value: str) -> bool:
return value.lower() in ("true", "yes", "1", "t", "y", "on")
+4 -13
View File
@@ -1,13 +1,12 @@
import base64
import logging
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.logging_utils import logging_error, logging_info
from amqp.adapter.serializer import map_as_string, parse_map_string
from amqp.adapter.trace_info_adapter import TraceInfoAdapter
from amqp.model.model import ServiceMessage
from amqp.model.service_message_type import ServiceMessageType
from amqp.model.snowflake_id import SnowflakeId
from amqp.adapter.serializer import parse_map_string, map_as_string
RS = "|"
SPLIT_RS = r"\|"
@@ -99,9 +98,7 @@ class ServiceMessageFactory:
i += 1
reply_to = input_str_array[i] if i < len(input_str_array) else ""
i += 1
trace_info = (
parse_map_string(input_str_array[i]) if i < len(input_str_array) else {}
)
trace_info = parse_map_string(input_str_array[i]) if i < len(input_str_array) else {}
i += 1
message = input_str_array[i] if i < len(input_str_array) else ""
return ServiceMessage(
@@ -111,11 +108,7 @@ class ServiceMessageFactory:
recipient_name=recipient_name,
reply_to=reply_to,
trace_info=trace_info,
message=(
base64.b64decode(message).decode()
if payload_type & 0x80
else message
),
message=(base64.b64decode(message).decode() if payload_type & 0x80 else message),
)
except Exception as e:
logging_error(
@@ -148,9 +141,7 @@ class ServiceMessageFactory:
:return: formatted message type
"""
num_type = (
message.message_type.value
if message.message_type
else ServiceMessageType.UNKNOWN.value
message.message_type.value if message.message_type else ServiceMessageType.UNKNOWN.value
) - 1
fmt = (num_type & 0x7F) | (message.payload_type & 0x80)
return f"{fmt:02X}"
+205
View File
@@ -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
+85 -17
View File
@@ -24,9 +24,7 @@ class AMQAdapter:
:param config: config parser to read configuration values
"""
self.generator_id = config.getint(
"CleverMicro-AMQ", self.PREFIX + "generator-id", fallback=164
)
self.generator_id = 0
self.service_name = config.get(
"CleverMicro-AMQ",
self.PREFIX + "service-name",
@@ -40,11 +38,72 @@ class AMQAdapter:
self.PREFIX + "require-authenticated-user",
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.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")
#
# Consul related configuration
self.consul_port = (
int(
os.getenv(
"CONSUL_PORT",
config.get("CleverMicro-AMQ", self.PREFIX + "consul-port", fallback="8500"),
)
)
if os.getenv("CONSUL_PORT") is not None
else config.getint("CleverMicro-AMQ", self.PREFIX + "consul-port", fallback=8500)
)
self.consul_host = os.getenv(
"CONSUL_HOST",
config.get("CleverMicro-AMQ", self.PREFIX + "consul-host", fallback="consul"),
)
self.consul_max_retries = (
int(
os.getenv(
"CONSUL_MAX_RETRIES",
config.get("CleverMicro-AMQ", self.PREFIX + "consul-max-retries", fallback="5"),
)
)
if os.getenv("CONSUL_MAX_RETRIES") is not None
else config.getint("CleverMicro-AMQ", self.PREFIX + "consul-max-retries", fallback=5)
)
self.consul_initial_retry_delay = (
float(
os.getenv(
"CONSUL_INITIAL_RETRY_DELAY",
config.get(
"CleverMicro-AMQ",
self.PREFIX + "consul-initial-retry-delay",
fallback="0.2",
),
)
)
if os.getenv("CONSUL_INITIAL_RETRY_DELAY") is not None
else config.getfloat(
"CleverMicro-AMQ", self.PREFIX + "consul-initial-retry-delay", fallback=0.2
)
)
self.consul_counter_key = (
os.getenv(
"CONSUL_COUNTER_KEY",
config.get(
"CleverMicro-AMQ",
self.PREFIX + "consul-counter-key",
fallback="service/ids/counter",
),
)
if os.getenv("CONSUL_COUNTER_KEY") is not None
else config.get(
"CleverMicro-AMQ",
self.PREFIX + "consul-counter-key",
fallback="service/ids/counter",
)
)
def adapter_prefix(self) -> str:
return f"{self.service_name}-{self.generator_id}-"
@@ -68,15 +127,9 @@ class Dispatch:
:param config: config parser to read configuration values
"""
self.use_dlq = config.getboolean(
"CleverMicro-AMQ", self.PREFIX + "use-dlq", fallback=False
)
self.amq_host = config.get(
"CleverMicro-AMQ", self.PREFIX + "amq-host", fallback="rabbitmq"
)
self.amq_port = config.getint(
"CleverMicro-AMQ", self.PREFIX + "amq-port", fallback=5672
)
self.use_dlq = config.getboolean("CleverMicro-AMQ", self.PREFIX + "use-dlq", fallback=False)
self.amq_host = config.get("CleverMicro-AMQ", self.PREFIX + "amq-host", fallback="rabbitmq")
self.amq_port = config.getint("CleverMicro-AMQ", self.PREFIX + "amq-port", fallback=5672)
self.amq_port_tls = config.getint(
"CleverMicro-AMQ", self.PREFIX + "amq-port-tls", fallback=5671
)
@@ -110,7 +163,7 @@ class Backpressure:
:param config: config parser to read configuration values
"""
self.threshold_threads = config.getint(
self.threshold_load = config.getint(
"CleverMicro-AMQ",
self.PREFIX + "threshold",
fallback=int(os.getenv("PARALLEL_WORKERS", 1)),
@@ -132,6 +185,10 @@ class Backpressure:
self.idle_duration = config.getint(
"CleverMicro-AMQ", self.PREFIX + "idle-duration", fallback=30
)
# Flag to enable/disable CPU monitoring
self.cpu_monitoring_enabled = config.getboolean(
"CleverMicro-AMQ", self.PREFIX + "cpu-monitoring-enabled", fallback=False
)
class AMQConfiguration:
@@ -146,7 +203,7 @@ class AMQConfiguration:
config.read(config_file_name)
else:
# Get the directory of the caller (the module where this function is called)
caller_frame = inspect.stack()[1] # 0 is current frame, 1 is caller
caller_frame = inspect.stack()[1] # 0 is current_availability frame, 1 is caller
caller_module = inspect.getmodule(caller_frame[0])
if caller_module:
caller_dir = os.path.dirname(os.path.abspath(caller_module.__file__))
@@ -171,3 +228,14 @@ class AMQConfiguration:
self.amq_adapter: AMQAdapter = AMQAdapter(config)
self.dispatch: Dispatch = Dispatch(config)
self.backpressure: Backpressure = Backpressure(config)
def get_unique_instance_id(self) -> int:
"""
Get a unique instance ID for this AMQ service.
Returns:
int: A unique instance ID
"""
if self.instance_id == -1:
logging_warning(f"Generated unique instance ID: {self.instance_id}")
return self.instance_id
+10 -2
View File
@@ -14,6 +14,12 @@ 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
# Consul configuration for CleverMicro AMQ Adapter
cm.amq-adapter.consul-host=consul
cm.amq-adapter.consul-port=8500
cm.amq-adapter.consul-max-retries=5
cm.amq-adapter.consul-initial-retry-delay=0.2
cm.amq-adapter.consul-id-generator-key=services/ids/counter
# CleverMicro AMQ Adapter dispatcher, settings related to RabbitMQ connection and optional REST forwarding
cm.dispatch.use-dlq=true
@@ -27,7 +33,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 +46,5 @@ 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
# AMQ Adapter can monitor CPU usage and trigger backpressure alerts based on CPU usage. Enable this feature explicitly
cm.backpressure.cpu-monitoring-enabled=false
-11
View File
@@ -1,11 +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)
+124 -69
View File
@@ -1,8 +1,10 @@
import base64
import json
from enum import Enum
from typing import List, Dict
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List
from aio_pika.abc import AbstractRobustConnection
from amqp.model.service_message_type import ServiceMessageType
from amqp.model.snowflake_id import SnowflakeId
@@ -16,9 +18,29 @@ Define the structure of messages used in Clever Micro
"""
class DataMessageMagicByte(Enum):
HTTP_REQUEST = "A"
RPC_REQUEST = "R"
@dataclass
class CleverMicroMessage:
DEFAULT_CONTENT_TYPE = ["application/octet-stream"]
"""
The common structure of a message that transfers the user request to the Service and back. All data related messages
conform to this base structure:
magic - identifies the message type (REST, RPC, DataResponse)
id - unique message ID
m_field - method name (for REST style calls) or service name (for RPC style calls) or response code for DataResponse
d_field - domain name (for REST style calls) or module name (for RPC style calls) or error code for DataResponse
p_field - path (for REST style calls) or method name (for RPC style calls) or error cause for DataResponse
headers - headers of the message (e.g. Content-Type, Accept, etc.)
trace_info - trace information (e.g. trace ID, span ID, etc. for Jaeger trace monitor)
base64body - the body of the message encoded in base64
"""
def __init__(
self,
magic: str,
@@ -66,11 +88,22 @@ class CleverMicroMessage:
def body(self) -> bytes:
return base64.b64decode(self.base64body())
def contentType(self) -> str:
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):
"""
The class represents a structure of a message that transfers the user request to the Service.
It may also be used to make a call/request between services.
The response to this call is in format of DataResponse class.
See DataMessageFactory for code to instantiate, serialize and deserialize the class
"""
def __init__(
self,
magic: str,
@@ -82,9 +115,7 @@ class DataMessage(CleverMicroMessage):
trace_info: Dict[str, str],
base64body: bytes,
):
super().__init__(
magic, id, method, domain, path, headers, trace_info, base64body
)
super().__init__(magic, id, method, domain, path, headers, trace_info, base64body)
def method(self) -> str:
return self.m_field()
@@ -95,19 +126,77 @@ class DataMessage(CleverMicroMessage):
def path(self) -> str:
return self.p_field()
def routing_key(self) -> str:
return sanitize_as_rabbitmq_name(f"{self.d_field()}.{self.p_field()}")
class MultipartDataMessage(DataMessage):
def __init__(self, message: DataMessage, extra_data: dict):
super().__init__(
message.magic(),
message.id(),
message.method(),
message.domain(),
message.path(),
message.headers(),
message.trace_info(),
message.base64body(),
)
self.extra_data = extra_data
@dataclass
class AMQMessage(DataMessage):
"""
Expand the core DataMessage type to include the channel, which is required to create a dedicated queue
in case the response should initiate file download.
"""
_connection: AbstractRobustConnection = field(repr=False) # Don't show in repr
def __init__(self, data_message: DataMessage, connection: AbstractRobustConnection):
# Initialize the parent DataMessage with all properties from the input data_message
super().__init__(
magic=data_message.magic(),
id=data_message.id(),
method=data_message.method(),
domain=data_message.domain(),
path=data_message.path(),
headers=data_message.headers(),
trace_info=data_message.trace_info(),
base64body=data_message.base64body(),
)
self._connection = connection
self.extra_data = data_message.extra_data if hasattr(data_message, "extra_data") else {}
@property
def connection(self) -> AbstractRobustConnection:
"""Get the AMQP connection associated with this message"""
return self._connection
class DataResponse(CleverMicroMessage):
"""
The response message to REST & RPC style calls (note: the request is made with DataMessage)
See DataResponseFactory for code to instantiate, serialize and deserialize the class
"""
def __init__(
self,
magic: str,
id: SnowflakeId,
response_code: str,
content_type: str,
error: str,
error_cause: str,
headers: Dict[str, List[str]],
trace_info: Dict[str, str],
base64body: bytes,
):
if headers is None:
headers = {}
headers["Content-Type"] = [content_type]
super().__init__(
magic,
id,
@@ -179,65 +268,6 @@ class AMQRoute:
return self.__str__()
@dataclass
class DataMessage:
"""
The class represents a structure of a message that transfers the user request to the Service.
It may also be used to make a call/request between services.
The response to this call is in format of DataResponse class.
Please see DataMessageFactory for code to instantiate, serialize and deserialize the class
"""
magic: str
id: SnowflakeId
method: str
domain: str
path: str
headers: dict[str, List[str]]
trace_info: dict[str, str]
base64_body: bytes
def routing_key(self) -> str:
return sanitize_as_rabbitmq_name(f"{self.domain}.{self.path}")
def body(self) -> bytes:
return base64.b64decode(self.base64_body)
class MultipartDataMessage(DataMessage):
def __init__(self, message: DataMessage, extra_data: dict):
super().__init__(
message.magic,
message.id,
message.method,
message.domain,
message.path,
message.headers,
message.trace_info,
message.base64_body,
)
self.extra_data = extra_data
@dataclass
class DataResponse:
"""
The response message to RPC style call (note: the request is made with DataMessage)
Please see DataResponseFactory for code to instantiate, serialize and deserialize the class
"""
id: str
response_code: int
content_type: str
response: bytes
error: str | None
error_cause: str | None
trace_info: dict[str, str]
def body(self) -> bytes:
return base64.b64decode(self.response)
# Ensure backward compatibility
AMQResponse = DataResponse
@@ -266,7 +296,7 @@ class ScalingRequestAlert(Enum):
IDLE = "IDLE"
UPDATE = "UPDATE"
SCALE_UP = "SCALE_UP"
SCSLE_DOWN = "SCALE_DOWN"
SCALE_DOWN = "SCALE_DOWN"
@dataclass(frozen=True)
@@ -301,7 +331,32 @@ class ScalingRequest:
task_id=data["taskId"],
max_availability=data["maxAvailability"],
current_availability=data["currentAvailability"],
request_type=ScalingRequestAlert(
data["requestType"]
), # Convert string to Enum
request_type=ScalingRequestAlert(data["requestType"]), # Convert string to Enum
)
class AMQErrorMessage:
"""
The error response message to RPC style call (note: the request is made with DataMessage)
"""
def __init__(
self,
response_code: int,
error: str,
detail: str,
content_type: str = "application/json",
):
self.response_code = response_code
self.error = error
self.detail = detail
self.content_type = content_type
def serialize(self) -> bytes:
return json.dumps(
{
"response_code": self.response_code,
"error": self.error,
"detail": self.detail,
}
).encode("utf-8")
+17 -28
View File
@@ -1,23 +1,22 @@
import logging
import aio_pika
from aio_pika.abc import AbstractRobustQueue, AbstractRobustExchange
from aio_pika.abc import AbstractRobustExchange, AbstractRobustQueue
from pika.exchange_type import ExchangeType
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.logging_utils import (
logging_error,
logging_debug,
logging_error,
logging_info,
logging_warning,
)
from amqp.model.model import DataMessage, AMQRoute
from amqp.model.model import AMQRoute, DataMessage
from amqp.router.router_base import (
AMQ_REPLY_TO_EXCHANGE,
DLQ_EXCHANGE,
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 = (
@@ -85,21 +84,15 @@ class RabbitMQClient:
After this the Service is listening on (consuming from) the queue for incoming messages.
"""
if self.router.is_open(self.channel):
queue_args = (
self.DLQ_ARGS if self.amq_configuration.dispatch.use_dlq else None
)
logging_info(
"RabbitMQClient register routes, cfg mapping: [%s]", route_mapping
)
queue_args = self.DLQ_ARGS if self.amq_configuration.dispatch.use_dlq else None
logging_info("RabbitMQClient register routes, cfg mapping: [%s]", route_mapping)
# convert the comma-separated string into list of AMQRoute objects
requested_routes = self.router.unwrap_route_list(route_mapping)
for route in requested_routes:
try:
# Note: AMQRoute allows one of the exchange or queue to be None, in which case the default is used
queue_name = (
route.queue
if route.queue
else self.router.get_consume_queue_name()
route.queue if route.queue else self.router.get_consume_queue_name()
)
logging_info("RabbitMQClient register route: [%s]", route)
queue: AbstractRobustQueue = await self.channel.declare_queue(
@@ -109,14 +102,10 @@ class RabbitMQClient:
auto_delete=False,
arguments=queue_args,
)
_exchange_name: str = (
route.exchange if route.exchange else AMQ_RPC_EXCHANGE
)
exchange: AbstractRobustExchange = (
await self.channel.declare_exchange(
_exchange_name: str = route.exchange if route.exchange else AMQ_RPC_EXCHANGE
exchange: AbstractRobustExchange = await self.channel.declare_exchange(
name=_exchange_name, type=ExchangeType.topic
)
)
await queue.bind(exchange, route.key)
_c_tag = await queue.consume(callback, no_ack=False)
# Register the route with the router and publish its detail to other adapters
@@ -200,16 +189,14 @@ class RabbitMQClient:
message
), # context path is a routing key
)
logging_info(
f"Msg Publish (no-reply), messageId: {message.id}, to: {route}"
)
logging_info(f"Msg Publish (no-reply), messageId: {message.id()}, to: {route}")
else:
# This is RPC call, there should be response
pika_message: aio_pika.message.Message = aio_pika.message.Message(
body=DataMessageFactory.serialize(message),
content_encoding="utf-8",
delivery_mode=2,
correlation_id=str(message.id),
correlation_id=str(message.id()),
reply_to=self.router.get_reply_to_queue_name(),
)
await self.rpc_exchange.publish(
@@ -220,12 +207,14 @@ class RabbitMQClient:
)
logging_info(
"Msg Publish (RPC call), messageId: %s, to: %s",
message.id.as_string(),
message.id().as_string(),
route.as_string(),
)
else:
logging_warning(
"Message not sent, no route for %s/%s", message.domain, message.path
"Message not sent, no route for %s/%s",
message.domain(),
message.path(),
)
except Exception as err:
+85
View File
@@ -0,0 +1,85 @@
import asyncio
import logging
from typing import Any, Dict
from amqp.config.amq_configuration import AMQConfiguration
from amqp.rabbitmq.user_management_service_client import (
UserManagementServiceClient,
UserManagementServiceException,
)
async def get_user_info_from_token(
client: UserManagementServiceClient, token: str
) -> Dict[str, Any]:
"""
Get user information from a JWT token.
Args:
client: User Management Service client
token: JWT token
Returns:
User information
Raises:
UserManagementServiceException: If an error occurs
"""
try:
# Validate the token
token_info = await client.validate_token(token)
# Extract user ID from token info
user_id = None
for item in token_info:
if isinstance(item, dict) and "sub" in item:
user_id = item["sub"]
break
if not user_id:
raise UserManagementServiceException(
"cleverthis.clevermicro.auth.invalid_token",
"Token does not contain user ID (sub claim)",
{},
)
# Query user information
user_info = await client.query_user_by_id(user_id)
return user_info
except UserManagementServiceException as e:
logging.error(f"User Management Service error: {e.exception_type}: {e.message}")
raise
except Exception as e:
logging.error(f"Unexpected error: {e}")
raise
async def main():
# Initialize configuration
config = AMQConfiguration("")
# Create client
client = UserManagementServiceClient(config)
try:
# Sample JWT token (replace with actual token)
token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
# Get user information
user_info = await get_user_info_from_token(client, token)
print(f"User information: {user_info}")
except UserManagementServiceException as e:
print(f"Error: {e.exception_type}: {e.message}")
if e.additional_info:
print(f"Additional info: {e.additional_info}")
except Exception as e:
print(f"Unexpected error: {e}")
finally:
# Close client
await client.close()
if __name__ == "__main__":
asyncio.run(main())
@@ -0,0 +1,260 @@
import asyncio
import json
import logging
import uuid
from typing import Any, Dict, List, Optional
import aio_pika
from aio_pika import Message
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustChannel, AbstractRobustConnection
from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
from amqp.config.amq_configuration import AMQConfiguration
class UserManagementServiceException(Exception):
"""Exception raised for errors in the User Management Service."""
def __init__(self, exception_type: str, message: str, additional_info: Dict[str, Any] = None):
self.exception_type = exception_type
self.message = message
self.additional_info = additional_info or {}
super().__init__(f"{exception_type}: {message}")
class UserManagementServiceClient:
"""
Client for the User Management Service that communicates over RabbitMQ.
Supports operations for token, user, group, tenant, and permission services.
"""
# Exchange name for User Management Service
EXCHANGE_NAME = "cleverthis.clevermicro.auth.endpoints"
# Service routing keys
TOKEN_SERVICE = "token-service-v1"
USER_SERVICE = "user-service-v1"
GROUP_SERVICE = "group-service-v1"
TENANT_SERVICE = "tenant-service-v1"
PERMISSION_SERVICE = "permission-service-v1"
# Response queue name
RESPONSE_QUEUE_NAME = "user-management-service-cleverswarm"
def __init__(self, configuration: AMQConfiguration, loop: asyncio.AbstractEventLoop = None):
"""
Initialize the User Management Service client.
Args:
configuration: AMQ configuration
loop: Event loop to use (optional)
"""
self.amq_configuration = configuration
self.loop = loop or asyncio.get_event_loop()
self.connection: Optional[AbstractRobustConnection] = None
self.channel: Optional[AbstractRobustChannel] = None
self.exchange: Optional[aio_pika.RobustExchange] = None
self.response_queue: Optional[aio_pika.RobustQueue] = None
self.response_consumer_tag: Optional[str] = None
self.futures: Dict[str, asyncio.Future] = {}
self.initialized = False
async def initialize(self) -> None:
"""Initialize the RabbitMQ connection and set up the exchange and response queue."""
if self.initialized:
return
try:
# Connect to RabbitMQ
self.connection = await aio_pika.connect_robust(
host=self.amq_configuration.dispatch.amq_host,
port=self.amq_configuration.dispatch.amq_port,
login=self.amq_configuration.dispatch.rabbit_mq_user,
password=self.amq_configuration.dispatch.rabbit_mq_password,
loop=self.loop,
)
# Create channel
self.channel = await self.connection.channel()
await self.channel.set_qos(prefetch_count=1)
# Declare exchange
self.exchange = await self.channel.declare_exchange(
name=self.EXCHANGE_NAME,
type=aio_pika.ExchangeType.TOPIC,
durable=True,
)
# Declare response queue
self.response_queue = await self.channel.declare_queue(
name=self.RESPONSE_QUEUE_NAME,
durable=True,
exclusive=False,
auto_delete=False,
)
# Start consuming responses
self.response_consumer_tag = await self.response_queue.consume(
self._on_response_callback,
no_ack=False,
)
self.initialized = True
logging_info("User Management Service client initialized")
except Exception as e:
logging_error(f"Failed to initialize User Management Service client: {e}")
if self.connection and not self.connection.is_closed:
await self.connection.close()
raise
async def _on_response_callback(self, message: AbstractIncomingMessage) -> None:
"""
Handle responses from the User Management Service.
Args:
message: The incoming message
"""
try:
correlation_id = message.correlation_id
if not correlation_id:
logging_error("Received response without correlation ID")
await message.ack()
return
future = self.futures.get(correlation_id)
if not future:
logging_error(f"No pending request found for correlation ID: {correlation_id}")
await message.ack()
return
try:
response_data = json.loads(message.body.decode('utf-8'))
logging_debug(f"Received response: {response_data}")
future.set_result(response_data)
except Exception as e:
future.set_exception(e)
await message.ack()
except Exception as e:
logging_error(f"Error processing response: {e}")
await message.ack()
async def close(self) -> None:
"""Close the RabbitMQ connection."""
if self.connection and not self.connection.is_closed:
await self.connection.close()
self.initialized = False
logging_info("User Management Service client closed")
async def call_service(
self, service: str, method: str, args: Dict[str, Any] = None
) -> Dict[str, Any]:
"""
Call a method on a User Management Service.
Args:
service: Service routing key (e.g., 'token-service-v1')
method: Method name to call
args: Method arguments
Returns:
Response data
Raises:
UserManagementServiceException: If the service returns an error
"""
if not self.initialized:
await self.initialize()
correlation_id = str(uuid.uuid4())
future = self.loop.create_future()
self.futures[correlation_id] = future
request_data = {
"method": method,
"args": args or {},
}
message = Message(
body=json.dumps(request_data).encode('utf-8'),
content_type='application/json',
correlation_id=correlation_id,
reply_to=self.RESPONSE_QUEUE_NAME,
)
try:
await self.exchange.publish(message, routing_key=service)
logging_info(f"Sent request to {service}.{method}: {args}")
# Wait for response
response = await future
# Process response
if response.get("type") == "OK":
return response.get("result", [])
else:
raise UserManagementServiceException(
response.get("exception", "unknown_error"),
response.get("message", "Unknown error"),
response.get("additional_info", {})
)
except asyncio.CancelledError:
raise
except UserManagementServiceException:
raise
except Exception as e:
logging_error(f"Error calling service {service}.{method}: {e}")
raise
finally:
self.futures.pop(correlation_id, None)
async def validate_token(self, token: str) -> Dict[str, Any]:
"""
Validate a JWT token.
Args:
token: JWT token to validate
Returns:
Token information
"""
return await self.call_service(
service=self.TOKEN_SERVICE,
method="validateToken",
args={"token": token}
)
async def query_user_by_id(self, user_id: str) -> Dict[str, Any]:
"""
Query user information by user ID.
Args:
user_id: User ID
Returns:
User information
"""
return await self.call_service(
service=self.USER_SERVICE,
method="queryUserById",
args={"userId": user_id}
)
async def list_users_in_group(self, group_id: str) -> List[Dict[str, Any]]:
"""
List users in a group.
Args:
group_id: Group ID
Returns:
List of users
"""
return await self.call_service(
service=self.GROUP_SERVICE,
method="listUsersInGroup",
args={"groupId": group_id}
)
+10 -8
View File
@@ -1,6 +1,5 @@
import logging
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
@@ -41,13 +40,9 @@ class RouteDatabase:
return re.compile(".*")
return re.compile(regex)
def matches(
self, routing_key_from_route: str, routing_key_from_message: str
) -> bool:
def matches(self, routing_key_from_route: str, routing_key_from_message: str) -> bool:
route_pattern = self.pattern(routing_key_from_route)
return bool(
route_pattern.match(sanitize_as_rabbitmq_name(routing_key_from_message))
)
return bool(route_pattern.match(sanitize_as_rabbitmq_name(routing_key_from_message)))
def wrap_routes(self) -> str:
return ",".join(str(route) for route in self.routes)
@@ -60,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)
+12 -20
View File
@@ -1,8 +1,7 @@
import logging
from typing import Callable, List, Set
from typing import Callable, List, Optional, Set
from amqp.adapter.amq_route_factory import AMQRouteFactory, NULL_ROUTE
from amqp.adapter.logging_utils import logging_error, logging_debug, logging_info
from amqp.adapter.amq_route_factory import NULL_ROUTE, AMQRouteFactory
from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
from amqp.model.model import AMQRoute, DataMessage
from amqp.router.route_database import RouteDatabase
from amqp.router.utils import sanitize_as_rabbitmq_name
@@ -28,7 +27,7 @@ class RouterBase:
self.route_removed_notifier: Callable[[AMQRoute], None] | None = None
def get_routing_key(self, message: DataMessage) -> str:
return sanitize_as_rabbitmq_name(f"{message.domain}.{message.path}")
return sanitize_as_rabbitmq_name(f"{message.domain()}.{message.path()}")
def wrap_routes(self) -> str:
return self.route_database.wrap_routes()
@@ -38,10 +37,7 @@ class RouterBase:
def unwrap_route_list(self, route_string: str) -> List[AMQRoute]:
logging_debug("RouterMaster.unwrapRoutes, route(s): %s", route_string)
routes = [
AMQRouteFactory.from_string(route_str)
for route_str in route_string.split(",")
]
routes = [AMQRouteFactory.from_string(route_str) for route_str in route_string.split(",")]
return [route for route in routes if route != NULL_ROUTE]
def get_routes(self) -> Set[AMQRoute]:
@@ -51,14 +47,15 @@ class RouterBase:
return self.route_database.find_route(domain, path)
def find_route_by_message(self, message: DataMessage) -> AMQRoute | None:
return self.route_database.find_route(message.domain, message.path)
return self.route_database.find_route(message.domain(), message.path())
def 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)
logging_info(
" [DBG] RouterMaster.addRoute [%d]: %s", len(amq_route_list), message
)
logging_info(" [DBG] RouterMaster.addRoute [%d]: %s", len(amq_route_list), message)
for route in amq_route_list:
logging_info("RouterMaster.addRoute, adding route: %s", route)
if not route.exchange and not route.queue and not route.key:
@@ -82,14 +79,9 @@ class RouterBase:
def remove_routes(self, message: str) -> None:
try:
to_remove = [
AMQRouteFactory.from_string(route_str)
for route_str in message.split(",")
]
to_remove = [AMQRouteFactory.from_string(route_str) for route_str in message.split(",")]
removed_count = sum(1 for route in to_remove if self.remove_route(route))
logging_info(
f"RouterMaster.removeRoutes(): {message}, removed={removed_count}"
)
logging_info(f"RouterMaster.removeRoutes(): {message}, removed={removed_count}")
except IOError as err:
logging_error(f"Can't remove route(s): {err}")
+15 -31
View File
@@ -1,26 +1,26 @@
import logging
from typing import Set
from aio_pika.abc import (
AbstractIncomingMessage,
AbstractRobustChannel,
AbstractRobustQueue,
AbstractIncomingMessage,
)
from amqp.adapter.amq_route_factory import AMQRouteFactory
from amqp.adapter.logging_utils import (
logging_debug,
logging_error,
logging_info,
logging_debug,
logging_warning,
)
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import ServiceMessage, AMQRoute
from amqp.model.model import AMQRoute, ServiceMessage
from amqp.model.service_message_type import ServiceMessageType
from amqp.router.router_base import (
ANY_RECIPIENT,
CM_C_REQ_QUEUE,
CM_REQUEST_EXCHANGE,
CM_RESPONSE_EXCHANGE,
ANY_RECIPIENT,
)
from amqp.router.router_producer import RouterProducer
@@ -77,17 +77,13 @@ class RouterConsumer(RouterProducer):
except Exception as ioe:
logging_error("RouterConsumer FATAL: cannot initialize, %s", ioe)
else:
logging_error(
"RouterConsumer FATAL: cannot initialize, RabbitMQ channel is not open."
)
logging_error("RouterConsumer FATAL: cannot initialize, RabbitMQ channel is not open.")
"""
* ServiceMessageType.ROUTING_DATA_REQUEST handler listening on queue 'CleverMicroRequest'
"""
async def handle_routing_data_request_message(
self, message: AbstractIncomingMessage
):
async def handle_routing_data_request_message(self, message: AbstractIncomingMessage):
# some debug statements, TODO - remove for prod release
delivery_tag = message.delivery_tag
debug = "Delivery.Properties = " + message.properties.__str__()
@@ -101,9 +97,7 @@ class RouterConsumer(RouterProducer):
await message.ack(False)
cm_message: ServiceMessage = self.service_message_factory.from_bytes(
message.body
)
cm_message: ServiceMessage = self.service_message_factory.from_bytes(message.body)
route_mapping: str = self.wrap_own_routes()
logging_debug(
"******[%s] (ROUTING_DATA_REQ) msg=%s, routeMap=[%s]",
@@ -124,9 +118,7 @@ class RouterConsumer(RouterProducer):
route_mapping,
cm_message.reply_to,
)
await self.publish_service_message(
service_message, self.cm_response_exchange
)
await self.publish_service_message(service_message, self.cm_response_exchange)
except Exception as err:
logging_error(
"******[%s] (ROUTING_DATA_REQ) Response to queue: %s failed: %s",
@@ -158,9 +150,7 @@ class RouterConsumer(RouterProducer):
serviceMessage: ServiceMessage = self.service_message_factory.of(
ServiceMessageType.ROUTING_DATA, route.as_string(), ANY_RECIPIENT
)
await self.publish_service_message(
serviceMessage, self.cm_response_exchange
)
await self.publish_service_message(serviceMessage, self.cm_response_exchange)
self.add_own_route(route)
logging_info("RouterConsumer registered Route: %s", route)
else:
@@ -182,16 +172,14 @@ class RouterConsumer(RouterProducer):
service_message: ServiceMessage = self.service_message_factory.of(
ServiceMessageType.ROUTING_DATA_REMOVE, route.as_string(), ANY_RECIPIENT
)
if not await self.publish_service_message(
service_message, self.cm_response_exchange
):
if not await self.publish_service_message(service_message, self.cm_response_exchange):
logging_warning(
"RouterConsumer couldn't remove current route: %s, channel not open.",
"RouterConsumer couldn't remove current_availability route: %s, channel not open.",
route,
)
except Exception as ioe:
logging_error("RouterConsumer failed remove current route: %s", ioe)
logging_error("RouterConsumer failed remove current_availability route: %s", ioe)
result: bool = self.remove_route(route)
return result
@@ -201,13 +189,9 @@ class RouterConsumer(RouterProducer):
* Cleanup - de-register routes with master.
"""
routes: Set[AMQRoute] = self.get_routes()
logging_info(
"RouterConsumer.cleanup() [%d]: %s", len(routes), self.wrap_routes()
)
logging_info("RouterConsumer.cleanup() [%d]: %s", len(routes), self.wrap_routes())
not_removed = sum(
1
for r in (self.remove_route_on_cleanup(route) for route in routes)
if not r
1 for r in (self.remove_route_on_cleanup(route) for route in routes) if not r
)
if not_removed > 0:
logging_warning(
+17 -30
View File
@@ -7,22 +7,20 @@
* forwards it to the application.
"""
import logging
from aio_pika import Message
from aio_pika.abc import (
AbstractIncomingMessage,
AbstractRobustChannel,
AbstractRobustQueue,
AbstractRobustConnection,
AbstractRobustExchange,
AbstractIncomingMessage,
AbstractRobustQueue,
)
from pika.exchange_type import ExchangeType
from amqp.adapter.logging_utils import (
logging_debug,
logging_error,
logging_info,
logging_debug,
logging_warning,
)
from amqp.adapter.service_message_factory import ServiceMessageFactory
@@ -30,13 +28,13 @@ from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import ServiceMessage
from amqp.model.service_message_type import ServiceMessageType
from amqp.router.router_base import (
RouterBase,
ANY_RECIPIENT,
CM_C_AMQ_QUEUE,
CM_P_REPLY_TO,
CM_P_RESP_QUEUE,
CM_REQUEST_EXCHANGE,
CM_RESPONSE_EXCHANGE,
CM_C_AMQ_QUEUE,
CM_P_RESP_QUEUE,
CM_P_REPLY_TO,
ANY_RECIPIENT,
RouterBase,
)
@@ -49,16 +47,14 @@ class RouterProducer(RouterBase):
"""
super().__init__()
self.connection: AbstractRobustConnection | None = None
self.service_message_factory = ServiceMessageFactory(
configuration.amq_adapter.service_name
)
self.service_message_factory = ServiceMessageFactory(configuration.amq_adapter.service_name)
self.amq_configuration: AMQConfiguration = configuration
self.channel: AbstractRobustChannel | None = None
self.cm_request_exchange: AbstractRobustExchange | None = None
self.cm_response_exchange: AbstractRobustExchange | None = None
logging_info(
f"RouterProducer.<init>: %s, route: %s",
"RouterProducer.<init>: %s, route: %s",
self.amq_configuration.amq_adapter.service_name,
self.amq_configuration.amq_adapter.route_mapping,
)
@@ -85,7 +81,7 @@ class RouterProducer(RouterBase):
name=CM_REQUEST_EXCHANGE, type=ExchangeType.fanout
)
#
# 2. declare ResponseExchange where to publish current RoutingData (response to
# 2. declare ResponseExchange where to publish current_availability RoutingData (response to
# a RoutingData request)
self.cm_response_exchange = await self.channel.declare_exchange(
name=CM_RESPONSE_EXCHANGE, type=ExchangeType.fanout
@@ -100,9 +96,7 @@ class RouterProducer(RouterBase):
arguments={},
)
# 2b. bind queue to CM_RESPONSE_EXCHANGE w/o any conditions (empty routing key)
await _response_queue.bind(
exchange=CM_RESPONSE_EXCHANGE, routing_key="#"
)
await _response_queue.bind(exchange=CM_RESPONSE_EXCHANGE, routing_key="#")
# 2c. listen for incoming RoutingData messages
_c_tag = await _response_queue.consume(
callback=self.handle_routing_data_message, no_ack=False
@@ -116,9 +110,7 @@ class RouterProducer(RouterBase):
logging_error("RouterProducer FATAL: cannot initialize, %s", ioe)
else:
logging_error(
"RouterProducer FATAL: cannot initialize, RabbitMQ channel is not open."
)
logging_error("RouterProducer FATAL: cannot initialize, RabbitMQ channel is not open.")
async def handle_routing_data_message(self, message: AbstractIncomingMessage):
"""
@@ -126,15 +118,13 @@ class RouterProducer(RouterBase):
* routing data, invokes addRoute or removeRoute depending on the type of the message.
"""
logging_info(
f"[%s] Received ROUTING DATA: [%s], msgId=%s",
"[%s] Received ROUTING DATA: [%s], msgId=%s",
message.delivery_tag,
message.body,
message.message_id,
)
await message.ack(multiple=False)
cm_message: ServiceMessage = self.service_message_factory.from_bytes(
message.body
)
cm_message: ServiceMessage = self.service_message_factory.from_bytes(message.body)
# ignore the message if it comes from ourselves (even if from different instance)
logging_debug(
"******[%s] ROUTING_DATA, msg=%s, routeMap=[%s] sn=%s, replyTo=%s",
@@ -146,8 +136,7 @@ class RouterProducer(RouterBase):
)
if (
cm_message.is_valid()
and not self.amq_configuration.amq_adapter.service_name
== cm_message.reply_to
and not self.amq_configuration.amq_adapter.service_name == cm_message.reply_to
):
if cm_message.message_type == ServiceMessageType.ROUTING_DATA:
self.add_routes(cm_message.message)
@@ -166,9 +155,7 @@ class RouterProducer(RouterBase):
serviceMessage: ServiceMessage = self.service_message_factory.of(
ServiceMessageType.ROUTING_DATA_REQUEST, "", ANY_RECIPIENT
)
if not await self.publish_service_message(
serviceMessage, self.cm_request_exchange
):
if not await self.publish_service_message(serviceMessage, self.cm_request_exchange):
logging_warning(
"RouterProducer could not request route data: %s channel is not open.",
CM_REQUEST_EXCHANGE,
+128 -56
View File
@@ -1,32 +1,33 @@
import asyncio
import os
import threading
import time
import os
from asyncio import Future, AbstractEventLoop
from asyncio import AbstractEventLoop, Future
from typing import Dict
from aio_pika import Message
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustExchange
from opentelemetry import trace
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from opentelemetry.trace import Tracer, TraceState
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from pika.exchange_type import ExchangeType
from amqp.adapter.backpressure_handler import BackpressureHandler
from amqp.adapter.cleverthis_service_adapter import AMQMessage, CleverThisServiceAdapter
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.data_response_factory import DataResponseFactory
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter, AMQMessage
from amqp.adapter.file_handler import StreamingFileHandler, FileHandler
from amqp.adapter.logging_utils import logging_error, logging_debug, logging_info
from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
from amqp.adapter.serializer import map_as_string
from amqp.adapter.service_message_factory import ServiceMessageFactory
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import (
DataResponse,
DataMessage,
DataMessageMagicByte,
DataResponse,
MultipartDataMessage,
)
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient
from amqp.adapter.serializer import map_as_string
def collect_trace_info():
@@ -75,21 +76,15 @@ class DataMessageHandler:
self.file_handler: FileHandler = file_handler
self.backpressure_handler: BackpressureHandler = backpressure_handler
async def inbound_data_message_callback_as_thread(
self, message: AbstractIncomingMessage
):
async def inbound_data_message_callback_as_thread(self, message: AbstractIncomingMessage):
"""
The callback method that is called when a DataMessage is received from RabbitMQ, starts new Thread.
"""
threading.Thread(
target=lambda: asyncio.run(
self.inbound_data_message_callback_no_thread(message)
)
target=lambda: asyncio.run(self.inbound_data_message_callback_no_thread(message))
).start()
async def inbound_data_message_callback_no_thread(
self, message: AbstractIncomingMessage
):
async def inbound_data_message_callback_no_thread(self, message: AbstractIncomingMessage):
"""
Receives the DataMessage as bytes from RabbitMQ queue and deserializes it to the DataMessage object
Ensures Jaeger trace-id propagation.
@@ -101,19 +96,17 @@ class DataMessageHandler:
amq_message = await self.reconstitute_data_message(message)
if amq_message is not None:
self.ensure_trace_info(amq_message)
if amq_message.magic == DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST:
if amq_message.magic() == DataMessageMagicByte.HTTP_REQUEST.value:
await self.run_http_request_magic(message, amq_message, self.loop)
else:
logging_error(
"Unknown MAGIC value: [%s], don't know how to interpret the message",
amq_message.magic,
amq_message.magic(),
)
asyncio.run_coroutine_threadsafe(message.nack(requeue=False), self.loop)
return
else:
logging_error(
"######[%s] No valid AMQ Message present.", message.delivery_tag
)
logging_error("######[%s] No valid AMQ Message present.", message.delivery_tag)
self.record_duration(start, message.delivery_tag)
asyncio.run_coroutine_threadsafe(message.ack(multiple=False), self.loop)
@@ -131,16 +124,16 @@ class DataMessageHandler:
:param loop: Event loop
:return: DataResponse
"""
# Increment the counter for parallel executions
self.backpressure_handler.increase_parallel_executions()
# Increment the counter for parallel executions and check resource availability
self.backpressure_handler.increase_current_load()
await self.backpressure_handler.check_overload_condition()
logging_info(
f"PARALLEL: Current Capacity [{self.backpressure_handler.current_parallel_executions}] of [{self.backpressure_handler.parallel_workers}]"
f"PARALLEL: Current Capacity [{self.backpressure_handler.current_availability}] of [{self.backpressure_handler.max_availability}]"
)
if message.reply_to:
self.reply_to = message.reply_to
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}")
@@ -163,11 +156,17 @@ class DataMessageHandler:
except Exception as e:
logging_error(f"Request handler: {e}")
self.backpressure_handler.decrease_parallel_executions()
self.backpressure_handler.decrease_current_load()
async def reconstitute_data_message(
self, message: AbstractIncomingMessage
) -> DataMessage | None:
"""
DataMessage may be serialized as string in separate queue (addressing == 1) or be a MultiPart message with
files streamed over individual queues, one per file.
This function uses 'addressing' mode from message headers to determine how to deserialize (reconstitute)
the original message.
"""
amq_message: DataMessage | None = None
delivery_tag = message.delivery_tag
addressing: int = message.headers.get("clevermicro.addressing", 0)
@@ -177,9 +176,7 @@ class DataMessageHandler:
amq_message = await DataMessageFactory.from_stream(
message.body, rabbitmq_url=self.rabbit_mq_url, loop=self.loop
)
logging_info(
f"######[{delivery_tag}] AMQ Message from stream: {amq_message}"
)
logging_info(f"######[{delivery_tag}] AMQ Message from stream: {amq_message}")
if amq_message is not None:
extra_data = await self.file_handler.on_message(
message=message,
@@ -188,26 +185,24 @@ class DataMessageHandler:
loop=self.loop,
output_dir=self.output_dir,
)
logging_info(
f"######[{delivery_tag}] Multipart download result: {extra_data}"
)
logging_info(f"######[{delivery_tag}] Multipart download result: {extra_data}")
amq_message = MultipartDataMessage(amq_message, extra_data)
if amq_message:
logging_info(
"######[%s] AMQ Message received, ctag=%s, msgId=%s, %s:%s, traceInfo=%s",
message.delivery_tag,
message.consumer_tag,
amq_message.id,
amq_message.method,
amq_message.path,
map_as_string(amq_message.trace_info),
amq_message.id(),
amq_message.method(),
amq_message.path(),
map_as_string(amq_message.trace_info()),
)
return amq_message
def ensure_trace_info(self, amq_message: DataMessage):
propagator = TraceContextTextMapPropagator()
context = propagator.extract(
carrier=amq_message.trace_info, context=trace.context_api.get_current()
carrier=amq_message.trace_info(), context=trace.context_api.get_current()
)
self.tracer.start_span(name="AMQ-Adapter-Consumer", context=context)
# logging_info("CTX=%s", context)
@@ -216,15 +211,15 @@ class DataMessageHandler:
# logging.info(" [*] CTX2=%s", context_with_span)
# propagator.inject(context_with_span, amq_message.trace_info, self.set_text)
# logging.info(" [***] NEW TI=%s CTX=%s SCOPE=%s",
# map_as_string(amq_message.trace_info), context_with_span, context)
# map_as_string(amq_message.trace_info()), context_with_span, context)
def record_duration(self, start, delivery_tag):
global last_data_message_time
last_data_message_time = time.time()
global last_backpressure_event_time
last_backpressure_event_time = time.time()
logging_info(
"######[%s] Done, single ACK: %sms.",
delivery_tag,
last_data_message_time - start,
last_backpressure_event_time - start,
)
async def send_reply(self, reply_to: str, response: DataResponse):
@@ -238,17 +233,15 @@ class DataMessageHandler:
if reply_to:
pika_message: Message = Message(
body=DataResponseFactory.serialize(response),
correlation_id=response.id,
correlation_id=response.id(),
content_type=(
response.content_type[0]
if isinstance(response.content_type, list)
else response.content_type
response.content_type()[0]
if isinstance(response.content_type(), list)
else response.content_type()
),
)
_res = asyncio.run_coroutine_threadsafe(
self.reply_to_exchange.publish(
message=pika_message, routing_key=reply_to
),
self.reply_to_exchange.publish(message=pika_message, routing_key=reply_to),
self.loop,
)
logging_debug(
@@ -260,21 +253,18 @@ class DataMessageHandler:
async def reply_received_callback(self, message: AbstractIncomingMessage):
"""
The reply for DataMessage that we sent out earlier - must match the received response with
the original request and respond to the caller.
The handler for reply for DataMessage that we sent out earlier -
the code matches the received response with the original request and passes the response to the caller.
:param message:
:return:
"""
reply_id = message.correlation_id
content_type = message.content_type
delivery_tag = message.delivery_tag
debug = f"######[{delivery_tag}] (RPC REPLY) env: {message.headers}, props: {message.properties}, BODY: {message.body.decode()}"
logging_debug(debug)
# TODO - figure out implementation of this part, as this is sending reply back either to a service,
# or to the user
logging_info(
f"OUTSTANDING RECV: {self.outstanding}, size={len(self.outstanding)}"
)
logging_info(f"OUTSTANDING RECV: {self.outstanding}, size={len(self.outstanding)}")
future: Future = self.outstanding.pop(reply_id, None)
# if reply is None:
@@ -293,3 +283,85 @@ class DataMessageHandler:
future.set_result(response.body())
await message.ack(multiple=False)
async def send_rpc(self, route, message: DataMessage) -> Future:
"""
Sends an RPC message to the specified route and returns a Future that will be completed
when the response is received.
:param route: AMQRoute object containing exchange and routing information
:param message: DataMessage to be sent
:return: Future that will be completed with the response payload
"""
# Create a Future to be completed when the response is received
future = asyncio.Future()
# Store the future in the outstanding dictionary using the message ID as the key
message_id = message.id().as_string()
self.outstanding[message_id] = future
# Create a Pika message with the correlation ID set to the message ID
pika_message = Message(
body=DataMessageFactory.serialize(message),
correlation_id=message_id,
reply_to=self.rabbit_mq_client.router.get_reply_to_queue_name(),
content_type=(
message.content_type()[0]
if isinstance(message.content_type(), list)
else message.content_type()
),
)
# Get the exchange from the route
rpc_exchange_name = route.exchange
exchange: AbstractRobustExchange = await self.rabbit_mq_client.channel.declare_exchange(
name=rpc_exchange_name, type=ExchangeType.topic
)
# Publish the message to the exchange with the component name as the routing key
await exchange.publish(message=pika_message, routing_key=route.component_name)
logging_debug(
"###### RPC Request Published to exchg:(%s) routing_key=%s, msgId=%s",
rpc_exchange_name,
route.component_name,
message_id,
)
return future
async def send_command(self, route, message: DataMessage) -> bool:
"""
Sends an RPC message to the specified route and returns a Future that will be completed
when the response is received.
:param route: AMQRoute object containing exchange and routing information
:param message: DataMessage to be sent
:return: Future that will be completed with the response payload
"""
# Create a Pika message with the correlation ID set to the message ID
pika_message = Message(
body=DataMessageFactory.serialize(message),
correlation_id=message.id().as_string(),
content_type=(
message.content_type()[0]
if isinstance(message.content_type(), list)
else message.content_type()
),
)
# Get the exchange from the route
rpc_exchange_name = route.exchange
exchange: AbstractRobustExchange = await self.rabbit_mq_client.channel.declare_exchange(
name=rpc_exchange_name, type=ExchangeType.topic
)
# Publish the message to the exchange with the component name as the routing key
await exchange.publish(message=pika_message, routing_key=route.component_name)
logging_debug(
"###### RPC Request Published to exchg:(%s) routing_key=%s, msgId=%s",
rpc_exchange_name,
route.component_name,
message.id().as_string(),
)
return True
+105 -20
View File
@@ -1,25 +1,25 @@
import asyncio
import logging
import logging.config
import os
import asyncio
from asyncio import Future
from threading import Thread
from typing import Optional
from aio_pika.abc import AbstractRobustExchange
from amqp.adapter.backpressure_handler import BackpressureHandler
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.amq_route_factory import AMQRouteFactory
from amqp.adapter.backpressure_handler import BackpressureHandler
from amqp.adapter.consul_global_id_generator import get_unique_instance_id
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()}")
@@ -38,24 +38,43 @@ class AMQService:
def __init__(self, amq_configuration: AMQConfiguration, service_adapter=None):
logging.info("***********************************************")
logging.info("********** AMQ Service *******************")
logging.info("************** AMQ Service ***************")
logging.info("***********************************************")
self.loop = asyncio.new_event_loop()
self.amq_configuration = amq_configuration
self.service_adapter = service_adapter
self.unique_id = get_unique_instance_id(amq_configuration.amq_adapter)
amq_configuration.amq_adapter.generator_id = self.unique_id
self.data_message_factory = DataMessageFactory.get_instance(
self.amq_configuration.amq_adapter.generator_id
self.unique_id & DataMessageFactory.MACHINE_ID_MASK
)
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.rabbit_mq_client = RabbitMQClient(self.amq_configuration)
self.tracer = initialize_trace(name=amq_configuration.amq_adapter.service_name)
self.amq_data_message_handler = None
self.service_message_factory = ServiceMessageFactory(
self.amq_configuration.amq_adapter.service_name
)
self.backpressure_thread: Optional[Thread] = None
self.backpressure_handler: Optional[BackpressureHandler] = None
self.maximum_availability = -1
# =======================================================================
# =======================================================================
# ============== P u b l i c 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))
@@ -65,11 +84,76 @@ class AMQService:
logging.info("***********************************************")
self.loop.run_forever()
async def init(self, loop, service_adapter):
_reply_to_exchange: AbstractRobustExchange = await self.rabbit_mq_client.init(
loop
def backpressure(self, current_availability: int, maximum: int = -1):
"""
Send Backpressure event to the Management Service. The Backpressure logic determines the type of the event
based on the current_availability and maximum values.
- OVERLOAD event is triggered immediately when the current_availability value is below threshold
typically set to 10% of the maximum value, or 0 if no maximum is provided.
- IDLE event is triggered when the current_availability value exeeds 80% of the maximum value or is greater
than 0 if no maximum is provided.
- UPDATE event is sent if neither OVERLOAD nor IDLE conditions are met, indicating that the backpressure is
within acceptable limits.
:param current_availability: Currently available capacity.
:param maximum: Maximum value of the backpressure metric. If undefined or set to lt 0 (e.g. -1),
the maximum is than determined from max value of current_availability seen over the time.
"""
if self.backpressure_handler is not None:
# watch maximum & current value to always have some valid maximum reference
if maximum < 0:
if current_availability > self.maximum_availability:
self.maximum_availability = current_availability
else:
if maximum > self.maximum_availability:
self.maximum_availability = maximum
asyncio.run_coroutine_threadsafe(
self.backpressure_handler.update_backpressure_value(
current_availability, maximum if maximum > 0 else self.maximum_availability
),
loop=self.backpressure_handler.loop,
)
backpressure_handler = BackpressureHandler(
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
def get_unique_instance_id(self) -> int:
"""
Provides cluster-wide unique instance ID for this AMQ service instance.
This ID is used to identify the service instance in the cluster and is unique across all instances.
IDs are managed by shared Consul KV store, ensuring that even if multiple instances of the service are running,
they will each have unique ID.
"""
return self.unique_id
# =======================================================================
# ====================== Internal Methods ============================
# =======================================================================
async def init(self, loop, service_adapter):
_reply_to_exchange: AbstractRobustExchange = await self.rabbit_mq_client.init(loop)
self.backpressure_handler = BackpressureHandler(
channel=self.rabbit_mq_client.channel,
loop=loop,
config=self.amq_configuration,
@@ -83,11 +167,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)
@@ -96,7 +181,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()
@@ -121,8 +206,8 @@ class AMQService:
)
return None
def send_message(self, message: DataMessage, future: Future):
self.amq_data_message_handler.outstanding[str(message.id)] = 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):
self.amq_data_message_handler.outstanding.pop(message_id, None)
+184 -90
View File
@@ -1,102 +1,196 @@
from opentelemetry import metrics, trace
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.trace import TracerProvider, Tracer
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
"""
OpenTelemetry configuration for tracing and metrics.
from prometheus_client import start_http_server
from opentelemetry.exporter.prometheus import PrometheusMetricReader
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
This module provides functionality to initialize and configure OpenTelemetry
for both tracing (Jaeger via OTLP) and metrics (Prometheus) in the AMQ adapter.
"""
import logging
import os
from dataclasses import dataclass
from typing import Any, Dict, Optional
from opentelemetry import metrics, trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.exporter.prometheus import PrometheusMetricReader
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from opentelemetry.sdk.resources import SERVICE_NAME, SERVICE_VERSION, Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
from opentelemetry.trace import Tracer
from 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")
)
@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,
}
)
#
# 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.
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
# Add OTLP exporter for Jaeger
otlp_exporter = setup_otlp_exporter(config)
provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
# Add console exporter in debug mode
if config.debug:
provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))
return provider
def setup_metrics_provider(config: MetricsConfig, resource: Resource) -> MeterProvider:
"""Set up the MeterProvider with Prometheus export."""
try:
# Start Prometheus HTTP server
start_http_server(port=config.prometheus_port, addr=config.prometheus_host)
logger.info(
f"Started Prometheus metrics server on {config.prometheus_host}:{config.prometheus_port}"
)
# Create Prometheus reader
prometheus_reader = PrometheusMetricReader()
# Create periodic reader for custom metrics
periodic_reader = PeriodicExportingMetricReader(
prometheus_reader, export_interval_millis=config.export_interval_millis
)
# Create and configure meter provider
return MeterProvider(resource=resource, metric_readers=[prometheus_reader, periodic_reader])
except Exception as e:
logger.error(f"Failed to set up metrics provider: {str(e)}")
# Return a no-op meter provider
return MeterProvider(resource=resource)
def initialize_telemetry(
tracing_config: Optional[TracingConfig] = None,
metrics_config: Optional[MetricsConfig] = None,
) -> Tracer:
"""
Initialize OpenTelemetry with tracing and metrics.
Args:
tracing_config: Configuration for tracing
metrics_config: Configuration for metrics
Returns:
Tracer: Configured OpenTelemetry tracer
Example:
```python
config = TracingConfig(
service_name="amq-adapter",
service_version="1.0.0",
otlp_endpoint="http://jaeger:4317",
debug=True
)
metrics_config = MetricsConfig(
prometheus_port=9464,
prometheus_host="0.0.0.0"
)
tracer = initialize_telemetry(config, metrics_config)
```
"""
# Use default configs if none provided
if tracing_config is None:
tracing_config = TracingConfig(
service_name=os.getenv("SERVICE_NAME", "amq-adapter"),
service_version=os.getenv("SERVICE_VERSION", "unknown"),
otlp_endpoint=os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317"),
insecure=os.getenv("OTEL_INSECURE_MODE", "true").lower() == "true",
debug=os.getenv("OTEL_DEBUG", "false").lower() == "true",
)
if metrics_config is None:
metrics_config = MetricsConfig(
prometheus_port=int(os.getenv("PROMETHEUS_PORT", "9464")),
prometheus_host=os.getenv("PROMETHEUS_HOST", "localhost"),
export_interval_millis=int(os.getenv("METRICS_EXPORT_INTERVAL", "30000")),
)
try:
# Create resource
resource = create_resource(tracing_config)
# Set up tracing
provider = setup_trace_provider(tracing_config, resource)
trace.set_tracer_provider(provider)
#
logger.info(f"Initialized tracing for service {tracing_config.service_name}")
# logging.info(' * Instrumenting application..... ')
# if Flask, then use:
# FlaskInstrumentor().instrument_app(flask_app)
# if FastAPI, then use
# FastAPIInstrumentor.instrument_app(app)
# Set up metrics
metrics_provider = setup_metrics_provider(metrics_config, resource)
metrics.set_meter_provider(metrics_provider)
logger.info("Initialized metrics export")
#
# 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.
# Return configured tracer
return trace.get_tracer(
tracing_config.service_name,
tracing_config.service_version,
)
# 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)
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")
return trace.get_tracer(name)
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 {},
)
+125
View File
@@ -0,0 +1,125 @@
"""
This module provides the CleverBRAG AMQP Adapter for handling API requests via AMQP.
The code IS MEANT to be part of CleverBRAG codebase, and has direct dependencies on the CleverBRAG codebase.
It is NOT MEANT to be used as a standalone module. It WILL NOT run outside the CleverBRAG context.
"""
from threading import Thread
from typing import Any, List
from uuid import UUID
# CleverBrag specific imports
import core
from fastapi.routing import APIRoute
from fastapi.security.http import HTTPAuthorizationCredentials
from pydantic import BaseModel
from shared.abstractions import R2RSerializable, User
# AMQP specific imports
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter
from amqp.config.amq_configuration import AMQConfiguration
from amqp.service.amq_service import AMQService
class BRAGArgument:
"""
This class is used to represent a BRAG API input argument. It is uses FastAPI annotation to define the arguments
(the argument type) that are passed as input arguments to the BRAG endpoint.
"""
def __init__(self, arg: Any):
self.arg = arg
self.module = arg.__module__
self.cls = arg.__name__ if hasattr(arg, "__name__") else str(arg)
self.is_r2r_serializable = issubclass(arg, R2RSerializable)
self.is_builtin = self.module == "builtins"
self.is_uuid = self.cls == "UUID"
self.is_optional = self.cls == "Optional"
self.is_list = self.cls == "list"
self.is_dict = self.cls == "dict"
self.is_base_model = issubclass(arg, BaseModel)
self.is_nested = self.is_list or self.is_optional or self.is_dict
if self.is_nested and hasattr(arg, "__args__"):
self.nested = BRAGArgument(getattr(arg, "__args__")[0])
if self.is_dict:
self.nested = (self.nested, BRAGArgument(getattr(arg, "__args__")[1]))
else:
self.nested = None
self.is_nested = False
def __repr__(self):
return f"{self.module}.{self.cls}" + (f"[{self.nested}]" if self.is_nested else "")
def convert_to_argument(arg: BRAGArgument, value: str) -> Any:
"""
Converts the argument to a string representation.
"""
try:
if arg.cls == "str":
return value
elif arg.cls == "bool":
return str(value).lower() == "true"
elif arg.cls == "float":
return float(value)
elif arg.cls == "int":
return int(value)
elif arg.is_uuid:
return UUID(value)
elif arg.is_optional:
return convert_to_argument(arg.nested, value) if arg.nested else None
elif arg.is_list:
return (
[convert_to_argument(arg.nested, v.strip()) for v in value.split(",")]
if arg.nested
else value.split()
)
elif arg.is_dict:
return (
{
convert_to_argument(arg.nested[0], k.strip()): convert_to_argument(
arg.nested[1], v.strip()
)
for k, v in (item.split(":") for item in value.split(","))
}
if arg.nested
else dict(item.split(":") for item in value.split(","))
)
elif arg.is_base_model:
return arg.arg.parse_raw(value)
else:
print(f"Unsupported argument type: {type(arg)}")
except Exception as e:
print(f"Error converting argument {arg}: {e}")
return None
class CleverBragAmqpAdapter(CleverThisServiceAdapter):
"""
CleverBRAG AMQP Adapter for CleverBRAG API. Implements the interface to CleverBRAG API.
"""
def __init__(
self,
amq_configuration: AMQConfiguration,
routes: List[APIRoute],
auth_provider: core.base.providers.auth.AuthProvider,
extra_init_data: dict,
):
super().__init__(
amq_configuration, routes, None, "auth_user", extra_init_data=extra_init_data
)
self.auth_provider = auth_provider
_amq_service: AMQService = AMQService(amq_configuration, self)
self.thread = Thread(target=_amq_service.run)
self.thread.start()
async def get_current_user(self, token: str) -> User:
"""
Overrides the default get_current_user method to use the token from the AMQP message.
:param token: JWT token from the AMQP message
:return: User object
"""
return await self.auth_provider.auth_wrapper(
HTTPAuthorizationCredentials(scheme="Bearer", credentials=token)
)
+44
View File
@@ -0,0 +1,44 @@
from clever_brag_adapter import BRAGArgument, CleverBragAmqpAdapter
from amqp.config.amq_configuration import AMQConfiguration
class R2RApp:
pass
def create_r2r_app():
return R2RApp()
if __name__ == "__main__":
#
# r2rapp: R2RApp = await create_r2r_app()
#
# In the standalone test context, need to start the loop as well. In the actual BRAG context,
# use 'await' instead.
r2rapp: R2RApp = create_r2r_app()
adapter: CleverBragAmqpAdapter = CleverBragAmqpAdapter(
AMQConfiguration("./application.properties.local"),
r2rapp.app.routes,
r2rapp.providers.auth,
{"app": r2rapp.config.app},
)
for method, routes in adapter.routes.items():
print(f"Method: {method}")
for _route in routes:
print(
f" Route: {_route.path} -> {_route.endpoint.__module__}.{_route.endpoint.__name__}"
)
_args = getattr(_route.endpoint, "__annotations__", {})
_verified_args = []
for arg in _args:
_arg = BRAGArgument(_args[arg])
_verified_args.append(_arg)
print(f" Arguments: {', '.join(map(str, _verified_args))}")
print("---------------------------")
print("Press ^C to exit...")
adapter.thread.join()
+149
View File
@@ -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
+34 -257
View File
@@ -5,129 +5,51 @@ It is NOT MEANT to be used as a standalone module. It WILL NOT run outside the C
"""
import importlib
import pathlib
import re
from collections import defaultdict
from threading import Thread
from typing import Optional, List, Dict
from aiohttp import ClientSession
from fastapi.routing import APIRoute
from fastapi.security import OAuth2PasswordRequestForm
from starlette.datastructures import UploadFile, Headers
from typing import Any, List, Optional
# CleverSwarm specific imports
import core.deps
from rest.v0 import api
from fastapi.routing import APIRoute
from schemas.user_schema import UserSchema
from amqp.adapter.data_parser import AMQDataParser
# AMQP specific imports
from amqp.adapter.data_response_factory import AMQResponseFactory
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter, AMQMessage
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQResponse
from amqp.service.amq_service import AMQService
# Suffix for the method name with serialization wrapper, prefer use this one over the orig endpoint
PREFERRED_SUFFIX = "_json"
def clone_and_adapt_route(original_route: APIRoute) -> APIRoute:
# =================================================================================================
# =========================== S u p p o r t i n g f u n c t i o n s =============================
# =================================================================================================
def get_callable_for_name(module_name: str, function_name: str) -> Optional[callable]:
"""
Creates a clone of an APIRoute. It checks if there's preferred endpoint with json serialization.
It also adjusts the path regex to ensure it matches the path in different API versions.
Get the function/Callable for the module_name.function_name if it exists in the module, else return None.
"""
# Check if the route has a preferred endpoint with json serialization
_preferred_endpoint = get_prefered_endpoint(
getattr(original_route.endpoint, "__module__"),
original_route.endpoint.__name__ + PREFERRED_SUFFIX,
)
# Need to adapt also the path regex to ensure it matches the path in different API versions
# Ensures the leading '^' is followed by any match and there's no '/' before final '$'
_new_regex_string = (
"^.*" + original_route.path_regex.pattern.lstrip("^").rstrip("/$") + "$"
)
_new_route: APIRoute = APIRoute(
path=original_route.path,
endpoint=(
_preferred_endpoint if _preferred_endpoint else original_route.endpoint
),
response_model=original_route.response_model,
status_code=original_route.status_code,
tags=original_route.tags.copy() if original_route.tags else None,
dependencies=(
original_route.dependencies.copy() if original_route.dependencies else None
),
summary=original_route.summary,
description=original_route.description,
response_description=original_route.response_description,
responses=original_route.responses.copy() if original_route.responses else None,
deprecated=original_route.deprecated,
name=original_route.name,
methods=original_route.methods.copy() if original_route.methods else None,
operation_id=original_route.operation_id,
response_model_include=original_route.response_model_include,
response_model_exclude=original_route.response_model_exclude,
response_model_by_alias=original_route.response_model_by_alias,
response_model_exclude_unset=original_route.response_model_exclude_unset,
response_model_exclude_defaults=original_route.response_model_exclude_defaults,
response_model_exclude_none=original_route.response_model_exclude_none,
include_in_schema=original_route.include_in_schema,
response_class=original_route.response_class,
dependency_overrides_provider=original_route.dependency_overrides_provider,
callbacks=original_route.callbacks.copy() if original_route.callbacks else None,
openapi_extra=(
original_route.openapi_extra.copy()
if original_route.openapi_extra
else None
),
generate_unique_id_function=original_route.generate_unique_id_function,
)
_new_route.path_regex = re.compile(_new_regex_string)
return _new_route
def get_prefered_endpoint(module_name: str, function_name: str) -> Optional[callable]:
"""Get the function if it exists in the module, else return None."""
try:
_module = importlib.import_module(module_name)
_prefered_endpoint = getattr(_module, function_name)
return _prefered_endpoint if callable(_prefered_endpoint) else None
_preferred_endpoint = getattr(_module, function_name)
return _preferred_endpoint if callable(_preferred_endpoint) else None
except Exception:
return None
def group_routes_by_method(api_routes: List[APIRoute]) -> Dict[str, List[APIRoute]]:
"""
Groups API routes by their HTTP method. This is important to ensure that the URLPath is matched
correctly, as different methods may share the same URLPath (e.g. GET, POST PUT, DELETE).
"""
_method_to_routes = defaultdict(
list
) # Automatically initializes new lists for new keys
for _route in api_routes:
if (
_route.methods
): # Ensure methods exist (should always be true for valid routes)
for method in _route.methods:
_method_to_routes[method].append(clone_and_adapt_route(_route))
return dict(_method_to_routes)
# =================================================================================================
# ======================== C l e v e r S w a r m A M Q A d a p t e r ========================
# =================================================================================================
class CleverSwarmAmqpAdapter(CleverThisServiceAdapter):
"""
CleverSwarm AMQP Adapter for CleverSwarm API. Implements the interface to CleverSwarm API.
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, session: ClientSession = None
):
super().__init__(amq_configuration, session)
self.routes = group_routes_by_method(api_routes=api.api_router.routes)
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()
@@ -135,170 +57,25 @@ class CleverSwarmAmqpAdapter(CleverThisServiceAdapter):
async def get_current_user(self, token: str) -> UserSchema:
"""
Overrides the default get_current_user method to use the token from the AMQP message.
In CleverSwarm context, it means simply to call provided get_current_user function.
:param token: JWT token from the AMQP message
:return: UserSchema object
"""
return await core.deps.get_current_user(token)
# ==================================================================================================
# ===================== C l e v e r S w a r m A P I m e t h o d s ============================
# ==================================================================================================
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.
async def get(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
"""
Handles the GET requests for the CleverSwarm API.
:param message: message from the AMQP that contains the GET request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: AMQResponse to be sent back to the API Gateway via AMQP
"""
return await self.handle_no_body_payload(message, auth_user, "GET")
:param endpoint: the Callable for the endpoint
async def put(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
:return: new Callable representing the JSON wrapper endpoint, if it exists or the original Callable.
"""
Handles the PUT requests for the CleverSwarm API.
:param message: message from the AMQP that contains the PUT request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: AMQResponse to be sent back to the API Gateway via AMQP
"""
return await self.handle_body_payload(message, auth_user, "PUT")
async def post(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
"""
Handles the POST requests for the CleverSwarm API.
:param message: message from the AMQP that contains the POST request, including body and headers
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: AMQResponse to be sent back to the API Gateway via AMQP
"""
return await self.handle_body_payload(message, auth_user, "POST")
async def head(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
"""
Handles the HEAD requests for the CleverSwarm API.
"""
return await self.handle_no_body_payload(message, auth_user, "HEAD")
async def delete(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
"""
Handles the DELETE requests for the CleverSwarm API.
"""
return await self.handle_no_body_payload(message, auth_user, "DELETE")
async def options(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
"""
Handles the OPTIONS requests for the CleverSwarm API.
"""
return await self.handle_no_body_payload(message, auth_user, "OPTIONS")
async def patch(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse:
"""
Handles the PATCH requests for the CleverSwarm API.
"""
return await self.handle_no_body_payload(message, auth_user, "PATCH")
async def handle_no_body_payload(
self, message: AMQMessage, auth_user: UserSchema, method: str
) -> AMQResponse:
"""
Handles the requests that do not have a body payload.
:param message: message from the AMQP that contains the request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: AMQResponse to be sent back to the API Gateway via AMQP
"""
_routes = self.routes.get(method, [])
_path, _args = AMQDataParser.parse_get_url(message.path)
for _route in _routes:
_match = _route.path_regex.match(_path)
if _match:
_args.update(_match.groupdict())
_args["authenticated_user"] = auth_user
# If the path matches, call the corresponding function
return await CleverThisServiceAdapter.call_cleverthis_api(
message,
_args,
api_method=_route.endpoint,
success_code=_route.status_code,
loop=self.loop,
)
return AMQResponseFactory.of(
message.id,
404,
"text/plain",
str(f"Destination location [{message.path}] does not exists").encode(
"utf-8"
),
message.trace_info,
)
async def handle_body_payload(
self, message: AMQMessage, auth_user: UserSchema, method: str
) -> AMQResponse:
"""
Handles the requests that have a body payload, including Multipart data.
:param message: message from the AMQP that contains the request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:param method: HTTP method (POST, PUT, etc.)
:return: AMQResponse to be sent back to the API Gateway via AMQP
"""
_form_data = AMQDataParser.parse_request_body(message)
_routes = self.routes.get(method, [])
_path, _query_args = AMQDataParser.parse_get_url(message.path)
for _route in _routes:
_match = _route.path_regex.match(_path)
if _match:
# If the path matches, call the corresponding function
_form_data.update(_match.groupdict())
_form_data.update(_query_args)
for _body_param in _route.dependant.body_params:
if _body_param.name in _form_data:
_file_data = _form_data[_body_param.name]
_form_data[_body_param.name] = (
_body_param.type_.__name__ == "UploadFile"
and message.extra_data.get(_body_param.name)
and UploadFile(
file=pathlib.Path(
message.extra_data.get(_body_param.name)
).open("rb"),
size=_file_data["size"],
filename=_file_data["filename"],
headers=Headers(
{"Content-Type": _file_data["mediaType"]}
),
)
or _form_data[_body_param.name]
)
_verified_args = {}
_required_args = getattr(_route.endpoint, "__annotations__", {})
for arg in _required_args:
if arg in _form_data:
_verified_args[arg] = _form_data[arg]
else:
# needed for /login input
if arg == "authenticated_user":
_verified_args["authenticated_user"] = auth_user
if _required_args[arg].__name__ == "OAuth2PasswordRequestForm":
try:
_verified_args[arg] = OAuth2PasswordRequestForm(
**_form_data
)
except Exception as e:
print(f"Error: {e}")
return await CleverThisServiceAdapter.call_cleverthis_api(
message,
_verified_args,
api_method=_route.endpoint,
success_code=_route.status_code if _route.status_code else 200,
loop=self.loop,
)
return AMQResponseFactory.of(
message.id,
404,
"text/plain",
str(f"Destination location [{message.path}] does not exists").encode(
"utf-8"
),
message.trace_info,
_preferred_endpoint = get_callable_for_name(
getattr(endpoint, "__module__"),
endpoint.__name__ + PREFERRED_SUFFIX,
)
return _preferred_endpoint if _preferred_endpoint is not None else endpoint
+2 -2
View File
@@ -1,7 +1,7 @@
from amqp.config.amq_configuration import AMQConfiguration
from clever_swarm_adapter import CleverSwarmAmqpAdapter
from amqp.config.amq_configuration import AMQConfiguration
if __name__ == "__main__":
adapter = CleverSwarmAmqpAdapter(AMQConfiguration("./application.properties.local"))
print("Press ^C to exit...")
+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_availability trace context is propagated
response = await client.post(
# IMPORTANT: Use the actual host and port where your FastAPI app is running
# If running directly on host, use http://localhost:8000
# 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
+20 -8
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "amq_adapter"
version = "0.2.27"
version = "0.2.34"
description = "The Python implementation of the AMQ Adapter for CleverMicro"
readme = "README.md"
authors = [
@@ -19,28 +19,40 @@ classifiers = [
]
requires-python = ">=3.11"
dependencies = [
"aio-pika~=9.5.5",
"aiohttp~=3.11.11",
"fastapi==0.114.2",
"aio-pika",
"aiohttp",
"fastapi",
"opentelemetry-api",
"opentelemetry-sdk",
"opentelemetry-exporter-otlp",
"opentelemetry-exporter-prometheus",
"opentelemetry-instrumentation-pika",
"pika~=1.3.2",
"pika",
"psutil",
"python_multipart"
"python-multipart",
"consul-kv"
]
[tool.setuptools.packages.find]
where = ["."]
include = ["amqp*"]
[tool.black]
line-length = 100
[tool.pytest.ini_options]
asyncio_default_fixture_loop_scope = "module"
[project.optional-dependencies]
dev = [
"pytest>=7.0",
"pytest-cov>=4.0",
"pytest-asyncio",
"build",
"twine"
"twine",
"black",
"isort",
"flake8",
"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,78 @@
import logging
import os
import threading
from amqp.adapter.consul_global_id_generator import get_unique_instance_id
from amqp.config.amq_configuration import AMQConfiguration
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(AMQConfiguration("").amq_adapter)
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}")
+205
View File
@@ -0,0 +1,205 @@
import asyncio
import time
import unittest
from unittest.mock import AsyncMock, MagicMock, patch
from amqp.adapter.backpressure_handler import (
BackpressureHandler,
ScaleRequestV1,
ScalingRequestAlert,
)
class TestBackpressureHandler(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.channel = AsyncMock()
self.loop = asyncio.get_event_loop()
self.config = MagicMock()
self.config.amq_adapter.swarm_service_id = "test-service"
self.config.amq_adapter.swarm_task_id = "test-task"
self.config.backpressure.threshold_load = 100
self.config.backpressure.time_window = 1
self.config.backpressure.idle_duration = 2
self.config.backpressure.cpu_monitoring_enabled = False
self.config.backpressure.cpu_overload_duration = 1
# Record callbacks for testing
BackpressureHandler._callback_list = {}
self.handler = BackpressureHandler(self.channel, self.loop, self.config)
self.handler.exchange = AsyncMock()
# Set do_loop to 1 to run the loop only once
self.handler.do_loop = 1
async def test_increase_decrease_current_load(self):
self.handler.current_load = 0
self.handler.increase_current_load()
self.assertEqual(self.handler.current_load, 1)
self.handler.decrease_current_load()
self.assertEqual(self.handler.current_load, 0)
async def test_update_current_load(self):
self.handler.current_availability = 0
self.handler.update_current_availability(5)
self.assertEqual(self.handler.current_availability, 5)
async def test_update_last_data_message_time(self):
old_time = self.handler.last_backpressure_event_time
self.handler.update_last_backpressure_event_time()
self.assertGreater(self.handler.last_backpressure_event_time, old_time)
@patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
async def test_update_backpressure_value(self, mock_publish):
# Test updating with a higher maximum
await self.handler.update_backpressure_value(50, 200)
self.assertEqual(self.handler.current_availability, 50)
self.assertEqual(self.handler.max_availability, 200)
# Test updating with a lower maximum (should not change)
await self.handler.update_backpressure_value(60, 50)
self.assertEqual(self.handler.current_availability, 60)
self.assertEqual(self.handler.max_availability, 200)
@patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
async def test_check_overload_condition_normal(self, mock_publish):
# Set up a normal state (60% available capacity)
self.handler.max_availability = 100
self.handler.current_availability = 60
self.handler.last_backpressure_event = ScalingRequestAlert.UPDATE
self.handler.last_backpressure_event_time = time.time() - 2 # 2 seconds ago
await self.handler.check_overload_condition()
# Should send an UPDATE event
mock_publish.assert_called_once()
args = mock_publish.call_args[0][0]
self.assertEqual(args.requestType, ScalingRequestAlert.UPDATE)
self.assertEqual(args.current_availability, 60)
self.assertEqual(args.max_availability, 100)
@patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
async def test_check_overload_condition_overload(self, mock_publish):
# Set up an overload state (10% available capacity)
self.handler.max_availability = 100
self.handler.current_availability = 10
self.handler.last_backpressure_event = ScalingRequestAlert.UPDATE
self.handler.last_backpressure_event_time = time.time() - 2 # 2 seconds ago
await self.handler.check_overload_condition()
# Should send an OVERLOAD event
mock_publish.assert_called_once()
args = mock_publish.call_args[0][0]
self.assertEqual(args.requestType, ScalingRequestAlert.OVERLOAD)
self.assertEqual(args.current_availability, 10)
self.assertEqual(args.max_availability, 100)
@patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
async def test_check_overload_condition_idle(self, mock_publish):
# Set up an idle state (90% available capacity)
self.handler.max_availability = 100
self.handler.current_availability = 90
self.handler.last_backpressure_event = ScalingRequestAlert.IDLE
# emulate fact that IDLE state is for 3 seconds longer than required minimum for sending IDLE event
self.handler.last_backpressure_event_time = (
time.time() - self.config.backpressure.idle_duration - 3
)
await self.handler.check_overload_condition()
# Should send an IDLE event
mock_publish.assert_called_once()
args = mock_publish.call_args[0][0]
self.assertEqual(args.requestType, ScalingRequestAlert.IDLE)
self.assertEqual(args.current_availability, 90)
self.assertEqual(args.max_availability, 100)
@patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
async def test_check_overload_condition_auto_max(self, mock_publish):
# Test with unset maximum that should be auto-detected
self.handler.max_availability = -1
self.handler.current_availability = 50
self.handler.last_backpressure_event = ScalingRequestAlert.UPDATE
self.handler.last_backpressure_event_time = (
time.time() - self.config.backpressure.time_window - 2
)
await self.handler.check_overload_condition()
# Should set maximum to current_availability and send UPDATE
# self.assertEqual(self.handler.max_availability, 50)
mock_publish.assert_called_once()
args = mock_publish.call_args[0][0]
self.assertEqual(args.requestType, ScalingRequestAlert.UPDATE)
self.assertEqual(args.current_availability, 50)
self.assertEqual(args.max_availability, -1)
# Test with a higher value that should update the maximum
mock_publish.reset_mock()
await self.handler.update_backpressure_value(80, 80)
# Should update maximum to 80
self.assertEqual(self.handler.max_availability, 80)
mock_publish.assert_called_once()
args = mock_publish.call_args[0][0]
self.assertEqual(args.current_availability, 80)
self.assertEqual(args.max_availability, 80)
@patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
async def test_handle_backpressure_overload_event(self, mock_publish):
self.handler.max_availability = 100
self.handler.current_availability = 10
await self.handler.handle_backpressure_overload_event()
mock_publish.assert_called_once()
args = mock_publish.call_args[0][0]
self.assertEqual(args.requestType, ScalingRequestAlert.OVERLOAD)
self.assertEqual(args.current_availability, 10)
self.assertEqual(args.max_availability, 100)
@patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
async def test_handle_backpressure_idle_event(self, mock_publish):
self.handler.max_availability = 100
self.handler.current_availability = 90
await self.handler._handle_backpressure_idle_event()
mock_publish.assert_called_once()
args = mock_publish.call_args[0][0]
self.assertEqual(args.requestType, ScalingRequestAlert.IDLE)
self.assertEqual(args.current_availability, 90)
self.assertEqual(args.max_availability, 100)
@patch("asyncio.run_coroutine_threadsafe")
async def test_publish_backpressure_request(self, mock_run):
scaling_request = ScaleRequestV1(
"test-service", "test-task", 100, 50, ScalingRequestAlert.UPDATE
)
await self.handler.publish_backpressure_request(scaling_request)
# Check that run_coroutine_threadsafe was called
self.assertEqual(mock_run.call_count, 1)
# Check that the callback was recorded
self.assertIn("_wrap_rabbit_mq_api", BackpressureHandler._callback_list)
class TestScaleRequestV1(unittest.TestCase):
def test_to_json(self):
request = ScaleRequestV1("test-service", "test-task", 100, 50, ScalingRequestAlert.UPDATE)
json_str = request.to_json()
# Check that the JSON contains the expected fields
self.assertIn('"version": 1', json_str)
self.assertIn('"serviceId": "test-service"', json_str)
self.assertIn('"taskId": "test-task"', json_str)
self.assertIn('"maxAvailability": 100', json_str)
self.assertIn('"currentAvailability": 50', json_str)
self.assertIn(f'"requestType": "{ScalingRequestAlert.UPDATE.value}"', json_str)
if __name__ == "__main__":
unittest.main()
+276 -72
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 CleverThisServiceAdapter, AMQMessage
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")
@@ -18,11 +109,13 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
# Create adapter instance
self.adapter = CleverThisServiceAdapter(
amq_configuration=self.mock_amq_config, session=self.mock_session
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)
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,16 +179,14 @@ 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:
with patch.object(self.adapter, "post", return_value=mock_response) as mock_post:
response = await self.adapter.on_message(mock_message)
self.assertEqual(response, mock_response)
mock_get_user.assert_not_called()
@@ -110,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
@@ -124,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
@@ -139,7 +232,7 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
response = await self.adapter.get(mock_message, None)
self.assertEqual(response.response_code, 200)
self.assertEqual(response.response_code(), "200")
self.mock_session.get.assert_called_once_with(
"http://testhost:8080/api/test",
headers={"X-Test": "value", "trace-id": "123"},
@@ -148,50 +241,50 @@ 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)
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
with patch.object(
self.adapter,
"handle_possible_form",
wraps=self.adapter.handle_possible_form,
"_authorize_forward_request",
wraps=self.adapter._authorize_forward_request,
) as mock_handler:
response = await self.adapter.put(mock_message, None)
self.assertEqual(response.response_code, 200)
self.assertEqual(response.response_code(), "200")
mock_handler.assert_called_once()
async def test_post_with_form_handling(self):
mock_message = 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"}')
@@ -199,28 +292,139 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
with patch.object(
self.adapter,
"handle_possible_form",
wraps=self.adapter.handle_possible_form,
"_authorize_forward_request",
wraps=self.adapter._authorize_forward_request,
) as mock_handler:
response = await self.adapter.post(mock_message, None)
self.assertEqual(response.response_code, 201)
self.assertEqual(response.response_code(), "201")
mock_handler.assert_called_once()
async def test_head_not_implemented(self):
mock_message = 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__":
@@ -0,0 +1,203 @@
import unittest
from unittest.mock import MagicMock, call, patch
from amqp.adapter.consul_global_id_generator import (
generate_fallback_id,
get_config_values,
get_unique_instance_id,
)
from amqp.config.amq_configuration import AMQConfiguration
class TestConsulGlobalIdGenerator(unittest.TestCase):
def setUp(self):
self.mock_config = MagicMock()
self.mock_config.consul_host = "test-consul"
self.mock_config.consul_port = 8500
self.mock_config.consul_counter_key = "test/counter/key"
self.mock_config.consul_max_retries = 3
self.mock_config.consul_initial_retry_delay = 0.1
def test_get_config_values_from_config(self):
"""Test retrieving configuration values from AMQAdapter config."""
host, port, key, retries, delay = get_config_values(self.mock_config)
self.assertEqual(host, "test-consul")
self.assertEqual(port, 8500)
self.assertEqual(key, "test/counter/key")
self.assertEqual(retries, 3)
self.assertEqual(delay, 0.1)
@patch.dict(
"os.environ",
{
"CONSUL_HOST": "env-consul",
"CONSUL_PORT": "8501",
"CONSUL_COUNTER_KEY": "env/counter/key",
"CONSUL_MAX_RETRIES": "4",
"CONSUL_RETRY_DELAY_SECONDS": "0.2",
},
)
def test_get_config_values_from_env(self):
"""Test retrieving configuration values from environment variables when config fails."""
self.mock_config = MagicMock(side_effect=Exception("Config error"))
host, port, key, retries, delay = get_config_values(AMQConfiguration("").amq_adapter)
self.assertEqual(host, "env-consul")
self.assertEqual(port, 8501)
self.assertEqual(key, "env/counter/key")
self.assertEqual(retries, 4)
self.assertEqual(delay, 0.2)
@patch("time.time", return_value=1000.0)
@patch("socket.gethostname", return_value="test-host")
@patch("random.randint", return_value=1234)
def test_generate_fallback_id(self, mock_randint, mock_hostname, mock_time):
"""Test fallback ID generation with controlled inputs."""
# Calculate expected value based on the mocked values
timestamp = int(1000.0 * 1000)
hostname_hash = hash("test-host") % 10000
random_part = 1234
expected_id = (timestamp * 100000000) + (hostname_hash * 10000) + random_part
result = generate_fallback_id()
self.assertEqual(result, expected_id)
@patch("consul_kv.Connection")
def test_get_unique_instance_id_success(self, mock_connection):
"""Test successful ID retrieval from Consul."""
mock_consul = MagicMock()
mock_connection.return_value = mock_consul
# Mock the get and put methods
mock_data = {"Value": b"42", "ModifyIndex": 123}
mock_consul.get.return_value = (0, mock_data)
mock_consul.put.return_value = True
result = get_unique_instance_id(self.mock_config)
# Should return the incremented value
self.assertEqual(result, 43)
mock_connection.assert_called_once_with(endpoint="test-consul:8500", timeout=5)
mock_consul.get.assert_called_once_with("test/counter/key")
mock_consul.put.assert_called_once_with("test/counter/key", "43", cas=123)
@patch("consul_kv.Connection")
def test_get_unique_instance_id_initialize(self, mock_connection):
"""Test initializing the counter when it doesn't exist."""
mock_consul = MagicMock()
mock_connection.return_value = mock_consul
# Return None to simulate counter not existing
mock_consul.get.return_value = (0, None)
mock_consul.put.return_value = True
result = get_unique_instance_id(self.mock_config)
# Should return 1 for the first ID
self.assertEqual(result, 1)
mock_consul.put.assert_called_once_with("test/counter/key", "1", cas=1)
@patch("consul_kv.Connection")
@patch("amqp.adapter.consul_global_id_generator.generate_fallback_id", return_value=9999)
def test_get_unique_instance_id_initialize_failure(self, mock_fallback, mock_connection):
"""Test fallback when initializing the counter fails."""
mock_consul = MagicMock()
mock_connection.return_value = mock_consul
# Return None to simulate counter not existing
mock_consul.get.return_value = (0, None)
mock_consul.put.return_value = False # Initialization fails
result = get_unique_instance_id(self.mock_config)
# Should return fallback ID
self.assertEqual(result, 9999)
mock_fallback.assert_called_once()
@patch("consul_kv.Connection")
@patch("time.sleep")
def test_get_unique_instance_id_retry_success(self, mock_sleep, mock_connection):
"""Test successful retry after CAS failure."""
mock_consul = MagicMock()
mock_connection.return_value = mock_consul
# First attempt fails, second succeeds
mock_data1 = {"Value": b"42", "ModifyIndex": 123}
mock_data2 = {"Value": b"43", "ModifyIndex": 124}
# First put fails, second succeeds
mock_consul.get.side_effect = [(0, mock_data1), (0, mock_data2)]
mock_consul.put.side_effect = [False, True]
result = get_unique_instance_id(self.mock_config)
# Should return the incremented value from the second attempt
self.assertEqual(result, 44)
self.assertEqual(mock_consul.get.call_count, 2)
self.assertEqual(mock_consul.put.call_count, 2)
mock_sleep.assert_called_once_with(0.1) # First retry delay
@patch("consul_kv.Connection")
@patch("time.sleep")
@patch("amqp.adapter.consul_global_id_generator.generate_fallback_id", return_value=9999)
def test_get_unique_instance_id_max_retries(self, mock_fallback, mock_sleep, mock_connection):
"""Test fallback after max retries."""
mock_consul = MagicMock()
mock_connection.return_value = mock_consul
# All attempts fail
mock_data = {"Value": b"42", "ModifyIndex": 123}
mock_consul.get.return_value = (0, mock_data)
mock_consul.put.return_value = False
result = get_unique_instance_id(self.mock_config)
# Should return fallback ID
self.assertEqual(result, 9999)
self.assertEqual(mock_consul.put.call_count, 3) # 3 retries as configured
mock_fallback.assert_called_once()
# Check exponential backoff
expected_calls = [
call(0.1), # First retry
call(0.2), # Second retry (doubled)
]
mock_sleep.assert_has_calls(expected_calls)
@patch("consul_kv.Connection")
@patch("amqp.adapter.consul_global_id_generator.generate_fallback_id", return_value=9999)
def test_get_unique_instance_id_consul_exception(self, mock_fallback, mock_connection):
"""Test fallback when Consul raises an exception."""
mock_connection.side_effect = Exception("Consul connection error")
result = get_unique_instance_id(self.mock_config)
# Should return fallback ID
self.assertEqual(result, 9999)
mock_fallback.assert_called_once()
@patch("consul_kv.Connection")
@patch("amqp.adapter.consul_global_id_generator.generate_fallback_id", return_value=9999)
def test_get_unique_instance_id_counter_disappeared(self, mock_fallback, mock_connection):
"""Test fallback when counter disappears during retry."""
mock_consul = MagicMock()
mock_connection.return_value = mock_consul
# First get succeeds, second returns None (counter disappeared)
mock_data = {"Value": b"42", "ModifyIndex": 123}
mock_consul.get.side_effect = [(0, mock_data), (0, None)]
mock_consul.put.return_value = False
result = get_unique_instance_id(self.mock_config)
# Should return fallback ID
self.assertEqual(result, 9999)
mock_fallback.assert_called_once()
if __name__ == "__main__":
unittest.main()
+13 -19
View File
@@ -10,25 +10,21 @@ class DataMessageFactoryTest(TestCase):
msg = DataMessageFactory.from_bytes(raw.encode("utf-8"))
self.assertIsNotNone(msg.id)
self.assertEqual(
msg.id.as_string(),
msg.id().as_string(),
"64F3B5AF4D00000A4000".lower(),
"SnowflakeID incorrectly parsed.",
)
self.assertEqual(
msg.headers["Content-Type"][0],
msg.headers()["Content-Type"][0],
"multipart/form-data; boundary=------------------------j1RjIqLrFH07XMsyMvpS1v",
)
def test_generate_snowflake_id(self):
_id = DataMessageFactory.get_instance(1).generate_snowflake_id()
self.assertEqual(
"00001000", _id.as_string()[-8:], "Instance ID not set correctly"
)
self.assertEqual("00001000", _id.as_string()[-8:], "Instance ID not set correctly")
_id = DataMessageFactory(2).generate_snowflake_id()
self.assertEqual(
"00002000", _id.as_string()[-8:], "Instance ID not set correctly"
)
self.assertEqual("00002000", _id.as_string()[-8:], "Instance ID not set correctly")
_f = DataMessageFactory(127)
_i0 = _f.generate_snowflake_id()
@@ -65,9 +61,9 @@ class DataMessageFactoryTest(TestCase):
{"k1": ["v1"], "k2": ["v2"]},
"Quick brown fox",
)
self.assertEqual("localhost", _req.domain, "Domain not set correctly")
self.assertEqual("/test/me", _req.path, "Context path not set correctly")
self.assertEqual("v1", _req.headers["k1"][0], "Headers not set correctly")
self.assertEqual("localhost", _req.domain(), "Domain not set correctly")
self.assertEqual("/test/me", _req.path(), "Context path not set correctly")
self.assertEqual("v1", _req.headers()["k1"][0], "Headers not set correctly")
self.assertEqual(b"Quick brown fox", _req.body(), "Body not correctly decoded")
def test_serialize(self):
@@ -82,11 +78,11 @@ class DataMessageFactoryTest(TestCase):
_ser = _f.serialize(_req)
self.assertEqual(65, _ser[2])
_deser = _f.from_bytes(_ser)
self.assertEqual(_req.id.as_string(), _deser.id.as_string())
self.assertEqual("localhost", _deser.domain)
self.assertEqual("/test/me", _deser.path)
self.assertEqual(_req.id().as_string(), _deser.id().as_string())
self.assertEqual("localhost", _deser.domain())
self.assertEqual("/test/me", _deser.path())
self.assertEqual(b"Quick brown fox", _deser.body())
self.assertEqual("v2", _deser.headers["k2"][0])
self.assertEqual("v2", _deser.headers()["k2"][0])
def test_amq_message_routing_key(self):
_f = DataMessageFactory(35)
@@ -111,7 +107,7 @@ class DataMessageFactoryTest(TestCase):
self.assertLessEqual(
_x_timestamp - _timestamp,
1,
"Timestamp in SnowflakeId differs too much from current time",
"Timestamp in SnowflakeId differs too much from current_availability time",
)
def test_to_string(self):
@@ -124,6 +120,4 @@ class DataMessageFactoryTest(TestCase):
"Quick brown fox",
)
_as_str = _f.to_string(_req)
self.assertGreater(
_as_str.index("domain='localhost'"), 0, "Expected content not found"
)
self.assertGreater(_as_str.index("domain='localhost'"), 0, "Expected content not found")
+165
View File
@@ -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")
+7 -6
View File
@@ -2,6 +2,7 @@ from unittest import TestCase
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.data_response_factory import DataResponseFactory
from amqp.model.model import DataMessageMagicByte
from amqp.model.snowflake_id import SnowflakeId
@@ -11,8 +12,8 @@ class TestDataResponseFactory(TestCase):
_f = DataResponseFactory()
_g = DataMessageFactory(111)
_resp = _f.create_async_response_message(_g.generate_snowflake_id())
self.assertEqual(200, _resp.response_code)
self.assertEqual(b"OK", _resp.response)
self.assertEqual("200", _resp.response_code())
self.assertEqual(b"OK", _resp.body())
def test_serialize(self):
_f = DataResponseFactory()
@@ -22,7 +23,7 @@ class TestDataResponseFactory(TestCase):
_ser = _f.serialize(_resp)
self.assertEqual(
2,
str(_ser[2:]).index(_f.MAGIC_BYTE_HTTP_REQUEST + _snowflakeId.as_string()),
str(_ser[2:]).index(DataMessageMagicByte.HTTP_REQUEST.value + _snowflakeId.as_string()),
"Magic byte and SnowflakeId not found",
)
self.assertGreater(str(_ser[2:]).index("|e=|"), 2)
@@ -35,6 +36,6 @@ class TestDataResponseFactory(TestCase):
_resp = _f.create_async_response_message(_snowflakeId)
_ser = _f.serialize(_resp)
_deser = _f.from_bytes(_ser)
self.assertEqual(_snowflakeId.as_string(), _deser.id)
self.assertEqual(200, _deser.response_code)
self.assertEqual(b"OK", _deser.response)
self.assertEqual(_snowflakeId.as_string(), _deser.id())
self.assertEqual("200", _deser.response_code())
self.assertEqual(b"OK", _deser.body())
+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()
View File
+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
+220
View File
@@ -0,0 +1,220 @@
import base64
import json
import unittest
from amqp.model.model import (
AMQErrorMessage,
AMQRoute,
CleverMicroMessage,
DataMessage,
DataResponse,
ScalingRequest,
ScalingRequestAlert,
)
from amqp.model.snowflake_id import SnowflakeId
class TestCleverMicroMessage(unittest.TestCase):
def setUp(self):
self.test_id = SnowflakeId()
self.test_headers = {"Content-Type": ["application/json"]}
self.test_trace_info = {"trace_id": "123"}
self.test_body = b"test data"
self.test_base64body = base64.b64encode(self.test_body)
self.message = CleverMicroMessage(
magic="A",
id=self.test_id,
m_field="GET",
d_field="test_domain",
p_field="test_path",
headers=self.test_headers,
trace_info=self.test_trace_info,
base64body=self.test_base64body,
)
def test_getters(self):
self.assertEqual(self.message.magic(), "A")
self.assertEqual(self.message.id(), self.test_id)
self.assertEqual(self.message.m_field(), "GET")
self.assertEqual(self.message.d_field(), "test_domain")
self.assertEqual(self.message.p_field(), "test_path")
self.assertEqual(self.message.headers(), self.test_headers)
self.assertEqual(self.message.trace_info(), self.test_trace_info)
self.assertEqual(self.message.base64body(), self.test_base64body)
def test_body_decoding(self):
self.assertEqual(self.message.body(), self.test_body)
def test_content_type(self):
self.assertEqual(self.message.content_type(), "application/json")
# Test default content type
message_no_content_type = CleverMicroMessage(
magic="A",
id=self.test_id,
m_field="GET",
d_field="test_domain",
p_field="test_path",
headers={},
trace_info={},
base64body=self.test_base64body,
)
self.assertEqual(
message_no_content_type.content_type(), CleverMicroMessage.DEFAULT_CONTENT_TYPE[0]
)
class TestDataMessage(unittest.TestCase):
def setUp(self):
self.test_id = SnowflakeId()
self.test_headers = {"Content-Type": ["application/json"]}
self.test_trace_info = {"trace_id": "123"}
self.test_base64body = base64.b64encode(b"test data")
self.message = DataMessage(
magic="A",
id=self.test_id,
method="GET",
domain="testdomain",
path="testpath",
headers=self.test_headers,
trace_info=self.test_trace_info,
base64body=self.test_base64body,
)
def test_getters(self):
self.assertEqual(self.message.method(), "GET")
self.assertEqual(self.message.domain(), "testdomain")
self.assertEqual(self.message.path(), "testpath")
def test_routing_key(self):
self.assertEqual(self.message.routing_key(), "testdomain.testpath")
class TestDataResponse(unittest.TestCase):
def setUp(self):
self.test_id = SnowflakeId()
self.test_trace_info = {"trace_id": "123"}
self.test_base64body = base64.b64encode(b"test data")
self.response = DataResponse(
magic="R",
id=self.test_id,
response_code="200",
content_type="application/json",
error="",
error_cause="",
headers={},
trace_info=self.test_trace_info,
base64body=self.test_base64body,
)
def test_getters(self):
self.assertEqual(self.response.response_code(), "200")
self.assertEqual(self.response.error(), "")
self.assertEqual(self.response.error_cause(), "")
self.assertEqual(self.response.content_type(), "application/json")
class TestAMQRoute(unittest.TestCase):
def setUp(self):
self.route = AMQRoute(
component_name="test_component",
exchange="test_exchange",
queue="test_queue",
key="test_key",
timeout=1000,
valid_until=1735689600000, # Some future timestamp
)
def test_str_representation(self):
expected = "test_component:test_exchange:test_queue:test_key:1000:1735689600000"
self.assertEqual(str(self.route), expected)
def test_equality(self):
same_route = AMQRoute(
component_name="test_component",
exchange="test_exchange",
queue="test_queue",
key="test_key",
timeout=1000,
valid_until=1735689600000,
)
different_route = AMQRoute(
component_name="other_component",
exchange="test_exchange",
queue="test_queue",
key="test_key",
timeout=1000,
valid_until=1735689600000,
)
self.assertEqual(self.route, same_route)
self.assertNotEqual(self.route, different_route)
class TestScalingRequest(unittest.TestCase):
def setUp(self):
self.request = ScalingRequest(
service_id="test_service",
task_id="test_task",
max_availability=5,
current_availability=3,
request_type=ScalingRequestAlert.OVERLOAD,
version=1,
)
def test_to_json(self):
json_str = self.request.to_json()
data = json.loads(json_str)
self.assertEqual(data["serviceId"], "test_service")
self.assertEqual(data["taskId"], "test_task")
self.assertEqual(data["maxAvailability"], 5)
self.assertEqual(data["currentAvailability"], 3)
self.assertEqual(data["requestType"], "OVERLOAD")
self.assertEqual(data["version"], 1)
def test_from_json(self):
json_str = json.dumps(
{
"serviceId": "test_service",
"taskId": "test_task",
"maxAvailability": 5,
"currentAvailability": 3,
"requestType": "OVERLOAD",
"version": 1,
}
)
request = ScalingRequest.from_json(json_str)
self.assertEqual(request.service_id, "test_service")
self.assertEqual(request.task_id, "test_task")
self.assertEqual(request.max_availability, 5)
self.assertEqual(request.current_availability, 3)
self.assertEqual(request.request_type, ScalingRequestAlert.OVERLOAD)
self.assertEqual(request.version, 1)
class TestAMQErrorMessage(unittest.TestCase):
def setUp(self):
self.error = AMQErrorMessage(
response_code=400,
error="Bad Request",
detail="Invalid input",
content_type="application/json",
)
def test_serialize(self):
serialized = self.error.serialize()
data = json.loads(serialized)
self.assertEqual(data["response_code"], 400)
self.assertEqual(data["error"], "Bad Request")
self.assertEqual(data["detail"], "Invalid input")
if __name__ == "__main__":
unittest.main()
+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")
@@ -0,0 +1,285 @@
import asyncio
import json
import unittest
from unittest.mock import AsyncMock, MagicMock, patch
from aio_pika import Message
from amqp.config.amq_configuration import AMQConfiguration
from amqp.rabbitmq.user_management_service_client import (
UserManagementServiceClient,
UserManagementServiceException
)
class TestUserManagementServiceClient(unittest.IsolatedAsyncioTestCase):
def setUp(self):
# Mock configuration
self.config = MagicMock(spec=AMQConfiguration)
self.config.dispatch = MagicMock()
self.config.dispatch.amq_host = "localhost"
self.config.dispatch.amq_port = 5672
self.config.dispatch.rabbit_mq_user = "guest"
self.config.dispatch.rabbit_mq_password = "guest"
# Create client
self.client = UserManagementServiceClient(self.config)
# Mock connection and channel
self.client.connection = MagicMock()
self.client.channel = MagicMock()
self.client.exchange = MagicMock()
self.client.exchange.publish = AsyncMock()
self.client.response_queue = MagicMock()
self.client.initialized = True
async def test_initialize(self):
# Mock aio_pika.connect_robust
with patch('aio_pika.connect_robust', new_callable=AsyncMock) as mock_connect:
# Mock connection
mock_connection = MagicMock()
mock_connect.return_value = mock_connection
# Mock channel - create an AsyncMock that returns the channel when awaited
mock_channel = AsyncMock()
mock_connection.channel = AsyncMock()
mock_connection.channel.return_value = mock_channel
mock_channel.set_qos = AsyncMock()
# Mock exchange
mock_exchange = MagicMock()
mock_channel.declare_exchange = AsyncMock(return_value=mock_exchange)
# Mock queue
mock_queue = MagicMock()
mock_channel.declare_queue = AsyncMock(return_value=mock_queue)
mock_queue.consume = AsyncMock(return_value="consumer-tag")
# Reset client
self.client.initialized = False
self.client.connection = None
self.client.channel = None
self.client.exchange = None
self.client.response_queue = None
# Call initialize
await self.client.initialize()
# Verify
mock_connect.assert_called_once_with(
host=self.config.dispatch.amq_host,
port=self.config.dispatch.amq_port,
login=self.config.dispatch.rabbit_mq_user,
password=self.config.dispatch.rabbit_mq_password,
loop=self.client.loop
)
mock_connection.channel.assert_called_once()
mock_channel.set_qos.assert_called_once_with(prefetch_count=1)
mock_channel.declare_exchange.assert_called_once()
mock_channel.declare_queue.assert_called_once()
mock_queue.consume.assert_called_once()
self.assertTrue(self.client.initialized)
async def test_call_service_success(self):
# Create a future for the response
future = asyncio.Future()
future.set_result({"type": "OK", "result": [{"name": "John Doe"}]})
# Mock the futures dictionary
with patch.dict(self.client.futures, {}, clear=True):
# Mock uuid.uuid4
with patch('uuid.uuid4', return_value="test-correlation-id"):
# Mock loop.create_future
with patch.object(self.client.loop, 'create_future', return_value=future):
# Call service
result = await self.client.call_service(
service=UserManagementServiceClient.TOKEN_SERVICE,
method="validateToken",
args={"token": "test-token"}
)
# Verify
self.client.exchange.publish.assert_called_once()
self.assertEqual(result, [{"name": "John Doe"}])
async def test_call_service_error(self):
# Create a future for the response
future = asyncio.Future()
future.set_result({
"type": "ERROR",
"exception": "cleverthis.clevermicro.auth.invalid_token",
"message": "Invalid token",
"additional_info": {"token": "test-token"}
})
# Mock the futures dictionary
with patch.dict(self.client.futures, {}, clear=True):
# Mock uuid.uuid4
with patch('uuid.uuid4', return_value="test-correlation-id"):
# Mock loop.create_future
with patch.object(self.client.loop, 'create_future', return_value=future):
# Call service and expect exception
with self.assertRaises(UserManagementServiceException) as context:
await self.client.call_service(
service=UserManagementServiceClient.TOKEN_SERVICE,
method="validateToken",
args={"token": "test-token"}
)
# Verify exception
exception = context.exception
self.assertEqual(exception.exception_type, "cleverthis.clevermicro.auth.invalid_token")
self.assertEqual(exception.message, "Invalid token")
self.assertEqual(exception.additional_info, {"token": "test-token"})
async def test_on_response_callback(self):
# Create a future
future = asyncio.Future()
# Add future to futures dictionary
self.client.futures["test-correlation-id"] = future
# Create a mock message
message = MagicMock()
message.correlation_id = "test-correlation-id"
message.body = json.dumps({"type": "OK", "result": [{"name": "John Doe"}]}).encode('utf-8')
message.ack = AsyncMock()
# Call callback
await self.client._on_response_callback(message)
# Verify
message.ack.assert_called_once()
self.assertTrue(future.done())
self.assertEqual(future.result(), {"type": "OK", "result": [{"name": "John Doe"}]})
async def test_on_response_callback_no_correlation_id(self):
# Create a mock message with no correlation ID
message = MagicMock()
message.correlation_id = None
message.ack = AsyncMock()
# Call callback
await self.client._on_response_callback(message)
# Verify
message.ack.assert_called_once()
async def test_on_response_callback_no_future(self):
# Create a mock message with unknown correlation ID
message = MagicMock()
message.correlation_id = "unknown-correlation-id"
message.ack = AsyncMock()
# Call callback
await self.client._on_response_callback(message)
# Verify
message.ack.assert_called_once()
async def test_validate_token(self):
# Mock call_service
with patch.object(self.client, 'call_service', new_callable=AsyncMock) as mock_call:
mock_call.return_value = [{"sub": "user-123", "name": "John Doe"}]
# Call validate_token
result = await self.client.validate_token("test-token")
# Verify
mock_call.assert_called_once_with(
service=UserManagementServiceClient.TOKEN_SERVICE,
method="validateToken",
args={"token": "test-token"}
)
self.assertEqual(result, [{"sub": "user-123", "name": "John Doe"}])
async def test_query_user_by_id(self):
# Mock call_service
with patch.object(self.client, 'call_service', new_callable=AsyncMock) as mock_call:
mock_call.return_value = [{"id": "user-123", "name": "John Doe", "email": "john@example.com"}]
# Call query_user_by_id
result = await self.client.query_user_by_id("user-123")
# Verify
mock_call.assert_called_once_with(
service=UserManagementServiceClient.USER_SERVICE,
method="queryUserById",
args={"userId": "user-123"}
)
self.assertEqual(result, [{"id": "user-123", "name": "John Doe", "email": "john@example.com"}])
async def test_close(self):
# Mock connection
self.client.connection.is_closed = False
self.client.connection.close = AsyncMock()
# Call close
await self.client.close()
# Verify
self.client.connection.close.assert_called_once()
self.assertFalse(self.client.initialized)
class TestGetUserInfoFromToken(unittest.IsolatedAsyncioTestCase):
async def test_get_user_info_from_token_success(self):
# Mock client
client = MagicMock(spec=UserManagementServiceClient)
client.validate_token = AsyncMock(return_value=[{"sub": "user-123"}])
client.query_user_by_id = AsyncMock(return_value=[{"id": "user-123", "name": "John Doe"}])
# Import the function
from amqp.rabbitmq.user_management_sample import get_user_info_from_token
# Call function
result = await get_user_info_from_token(client, "test-token")
# Verify
client.validate_token.assert_called_once_with("test-token")
client.query_user_by_id.assert_called_once_with("user-123")
self.assertEqual(result, [{"id": "user-123", "name": "John Doe"}])
async def test_get_user_info_from_token_no_user_id(self):
# Mock client
client = MagicMock(spec=UserManagementServiceClient)
client.validate_token = AsyncMock(return_value=[{"name": "John Doe"}]) # No sub claim
# Import the function
from amqp.rabbitmq.user_management_sample import get_user_info_from_token
# Call function and expect exception
with self.assertRaises(UserManagementServiceException) as context:
await get_user_info_from_token(client, "test-token")
# Verify exception
exception = context.exception
self.assertEqual(exception.exception_type, "cleverthis.clevermicro.auth.invalid_token")
self.assertEqual(exception.message, "Token does not contain user ID (sub claim)")
async def test_get_user_info_from_token_service_exception(self):
# Mock client
client = MagicMock(spec=UserManagementServiceClient)
client.validate_token = AsyncMock(
side_effect=UserManagementServiceException(
"cleverthis.clevermicro.auth.invalid_token",
"Invalid token",
{}
)
)
# Import the function
from amqp.rabbitmq.user_management_sample import get_user_info_from_token
# Call function and expect exception
with self.assertRaises(UserManagementServiceException) as context:
await get_user_info_from_token(client, "test-token")
# Verify exception
exception = context.exception
self.assertEqual(exception.exception_type, "cleverthis.clevermicro.auth.invalid_token")
self.assertEqual(exception.message, "Invalid token")
if __name__ == "__main__":
unittest.main()
+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
+2 -2
View File
@@ -2,8 +2,8 @@ from unittest import TestCase
import pytest
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.amq_route_factory import AMQRouteFactory
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.model.model import AMQRoute
from amqp.router.route_database import RouteDatabase
@@ -80,7 +80,7 @@ class TestRouteDatabase(TestCase):
message = DataMessageFactory.get_instance(111).create_request_message(
"POST", "clever3-book.localhost", "/aa/bb?cc=d", {}, ""
)
route = self._database.find_route(message.domain, message.path)
route = self._database.find_route(message.domain(), message.path())
self.assertIsNotNone(route)
self.assertEqual(ROUTE_3, route.as_string())
route = self._database.find_route("clever3.book.localhost", "/three-hot-dogs")
+210
View File
@@ -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)
View File
+419
View File
@@ -0,0 +1,419 @@
import asyncio
import unittest
from unittest.mock import AsyncMock, MagicMock, patch
from amqp.adapter.backpressure_handler import BackpressureHandler
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.file_handler import FileHandler
from amqp.adapter.service_message_factory import ServiceMessageFactory
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import (
DataMessage,
DataMessageMagicByte,
DataResponse,
MultipartDataMessage,
)
from amqp.model.snowflake_id import SnowflakeId
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient
from amqp.service.amq_message_handler import (
DataMessageHandler,
collect_trace_info,
get_default_trace_state,
set_text,
)
class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
def setUp(self):
# Mock dependencies
self.service_adapter = MagicMock(spec=CleverThisServiceAdapter)
self.service_adapter.on_message = AsyncMock()
self.tracer = MagicMock()
self.tracer.start_span = MagicMock()
self.message_factory = MagicMock(spec=DataMessageFactory)
self.service_message_factory = MagicMock(spec=ServiceMessageFactory)
self.reply_to_exchange = MagicMock()
self.reply_to_exchange.publish = AsyncMock()
self.reply_to_exchange.name = "test-exchange"
self.rabbit_mq_client = MagicMock(spec=RabbitMQClient)
self.rabbit_mq_client.connection = MagicMock()
self.amq_configuration = MagicMock(spec=AMQConfiguration)
self.amq_configuration.dispatch = MagicMock()
self.amq_configuration.dispatch.rabbit_mq_url = "amqp://guest:guest@localhost:5672/"
self.amq_configuration.dispatch.download_dir = "/tmp/test_downloads"
self.loop = asyncio.get_event_loop()
self.backpressure_handler = MagicMock(spec=BackpressureHandler)
self.backpressure_handler.increase_current_load = MagicMock()
self.backpressure_handler.decrease_current_load = MagicMock()
self.backpressure_handler.check_overload_condition = AsyncMock()
self.backpressure_handler.current_availability = 5
self.backpressure_handler.max_availability = 10
self.file_handler = MagicMock(spec=FileHandler)
self.file_handler.on_message = AsyncMock(return_value={})
# Create the handler
self.handler = DataMessageHandler(
service_adapter=self.service_adapter,
tracer=self.tracer,
message_factory=self.message_factory,
service_message_factory=self.service_message_factory,
reply_to_exchange=self.reply_to_exchange,
rabbit_mq_client=self.rabbit_mq_client,
amq_configuration=self.amq_configuration,
loop=self.loop,
backpressure_handler=self.backpressure_handler,
file_handler=self.file_handler,
)
# Create a mock message
self.mock_message = MagicMock()
self.mock_message.delivery_tag = 123
self.mock_message.consumer_tag = "consumer-tag"
self.mock_message.reply_to = "reply-queue"
self.mock_message.correlation_id = "correlation-id"
self.mock_message.headers = {"clevermicro.addressing": 0}
self.mock_message.properties = {}
self.mock_message.body = b"test-body"
self.mock_message.ack = AsyncMock()
self.mock_message.nack = AsyncMock()
# Create a mock data message
self.mock_data_message = MagicMock(spec=DataMessage)
self.mock_data_message.magic.return_value = DataMessageMagicByte.HTTP_REQUEST.value
self.mock_data_message.id.return_value = SnowflakeId.from_hex("1234567890abcdef1234")
self.mock_data_message.method.return_value = "GET"
self.mock_data_message.path.return_value = "/test/path"
self.mock_data_message.trace_info.return_value = {
"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
}
self.mock_data_message.body.return_value = b'{"test": "data"}'
# Create a mock data response
self.mock_data_response = MagicMock(spec=DataResponse)
self.mock_data_response.id.return_value = "response-id"
self.mock_data_response.content_type.return_value = "application/json"
# Patch DataMessageFactory.from_bytes
self.from_bytes_patch = patch(
"amqp.adapter.data_message_factory.DataMessageFactory.from_bytes",
return_value=self.mock_data_message,
)
self.mock_from_bytes = self.from_bytes_patch.start()
# Patch DataMessageFactory.from_stream
self.from_stream_patch = patch(
"amqp.adapter.data_message_factory.DataMessageFactory.from_stream",
return_value=self.mock_data_message,
)
self.mock_from_stream = self.from_stream_patch.start()
# Patch DataResponseFactory.serialize
self.serialize_patch = patch(
"amqp.adapter.data_response_factory.DataResponseFactory.serialize",
return_value=b"serialized-response",
)
self.mock_serialize = self.serialize_patch.start()
# Patch DataResponseFactory.from_bytes
self.response_from_bytes_patch = patch(
"amqp.adapter.data_response_factory.DataResponseFactory.from_bytes",
return_value=self.mock_data_response,
)
self.mock_response_from_bytes = self.response_from_bytes_patch.start()
# Patch asyncio.run_coroutine_threadsafe
self.run_coroutine_patch = patch("asyncio.run_coroutine_threadsafe")
self.mock_run_coroutine = self.run_coroutine_patch.start()
# Patch time.time
self.time_patch = patch("time.time", return_value=1234567890.0)
self.mock_time = self.time_patch.start()
# Patch os.remove
self.os_remove_patch = patch("os.remove")
self.mock_os_remove = self.os_remove_patch.start()
def tearDown(self):
self.from_bytes_patch.stop()
self.from_stream_patch.stop()
self.serialize_patch.stop()
self.response_from_bytes_patch.stop()
self.run_coroutine_patch.stop()
self.time_patch.stop()
self.os_remove_patch.stop()
async def test_inbound_data_message_callback_no_thread_success(self):
"""Test successful processing of an inbound message."""
# Setup
self.service_adapter.on_message.return_value = self.mock_data_response
# Execute
await self.handler.inbound_data_message_callback_no_thread(self.mock_message)
# Verify
self.mock_from_bytes.assert_called_once_with(self.mock_message.body)
self.tracer.start_span.assert_called_once()
self.backpressure_handler.increase_current_load.assert_called_once()
self.service_adapter.on_message.assert_called_once()
self.reply_to_exchange.publish.assert_called_once()
self.backpressure_handler.decrease_current_load.assert_called_once()
self.mock_message.ack.assert_called_once_with(multiple=False)
async def test_inbound_data_message_callback_no_thread_unknown_magic(self):
"""Test handling of a message with unknown magic value."""
# Setup
self.mock_data_message.magic.return_value = "UNKNOWN"
# Execute
await self.handler.inbound_data_message_callback_no_thread(self.mock_message)
# Verify
self.mock_message.nack.assert_called_once_with(requeue=False)
# self.mock_message.ack.assert_called_once_with(multiple=False)
async def test_inbound_data_message_callback_no_thread_no_message(self):
"""Test handling when no valid message is present."""
# Setup
self.mock_from_bytes.return_value = None
# Execute
await self.handler.inbound_data_message_callback_no_thread(self.mock_message)
# Verify
self.mock_message.ack.assert_called_once_with(multiple=False)
async def test_inbound_data_message_callback_as_thread(self):
"""Test the thread-based callback method."""
# Setup
with patch("threading.Thread") as mock_thread:
mock_thread_instance = MagicMock()
mock_thread.return_value = mock_thread_instance
# Execute
await self.handler.inbound_data_message_callback_as_thread(self.mock_message)
# Verify
mock_thread.assert_called_once()
mock_thread_instance.start.assert_called_once()
async def test_reconstitute_data_message_standard(self):
"""Test reconstituting a standard data message."""
# Execute
result = await self.handler.reconstitute_data_message(self.mock_message)
# Verify
self.assertEqual(result, self.mock_data_message)
self.mock_from_bytes.assert_called_once_with(self.mock_message.body)
async def test_reconstitute_data_message_stream(self):
"""Test reconstituting a message from a stream."""
# Setup
self.mock_message.headers = {"clevermicro.addressing": 1}
# Execute
result = await self.handler.reconstitute_data_message(self.mock_message)
# Verify
self.assertEqual(MultipartDataMessage, result.__class__)
self.mock_from_stream.assert_called_once()
self.file_handler.on_message.assert_called_once()
async def test_run_http_request_magic_success(self):
"""Test successful processing of an HTTP request."""
# Setup
self.service_adapter.on_message.return_value = self.mock_data_response
# Execute
await self.handler.run_http_request_magic(
self.mock_message, self.mock_data_message, self.loop
)
# Verify
self.backpressure_handler.increase_current_load.assert_called_once()
self.backpressure_handler.check_overload_condition.assert_called()
self.service_adapter.on_message.assert_called_once()
self.reply_to_exchange.publish.assert_called_once()
self.backpressure_handler.decrease_current_load.assert_called_once()
async def test_run_http_request_magic_multipart(self):
"""Test processing a multipart message with file cleanup."""
# Setup
self.service_adapter.on_message.return_value = self.mock_data_response
multipart_message = MultipartDataMessage(
self.mock_data_message, {"file1": "/tmp/file1.txt"}
)
# Execute
await self.handler.run_http_request_magic(self.mock_message, multipart_message, self.loop)
# Verify
self.mock_os_remove.assert_called_once_with("/tmp/file1.txt")
async def test_run_http_request_magic_service_exception(self):
"""Test handling of an exception from the service adapter."""
# Setup
self.service_adapter.on_message.side_effect = Exception("Service error")
# Execute
await self.handler.run_http_request_magic(
self.mock_message, self.mock_data_message, self.loop
)
# Verify
self.backpressure_handler.decrease_current_load.assert_called_once()
self.backpressure_handler.check_overload_condition.assert_called()
async def test_run_http_request_magic_reply_exception(self):
"""Test handling of an exception when sending the reply."""
# Setup
self.service_adapter.on_message.return_value = self.mock_data_response
self.reply_to_exchange.publish.side_effect = Exception("Reply error")
# Execute
await self.handler.run_http_request_magic(
self.mock_message, self.mock_data_message, self.loop
)
# Verify
# [no ACK at this level, ACK is in parent call] self.mock_message.ack.assert_called_once_with(requeue=False)
self.backpressure_handler.decrease_current_load.assert_called_once()
async def test_send_reply_success(self):
"""Test successful sending of a reply."""
# Execute
await self.handler.send_reply("reply-queue", self.mock_data_response)
# Verify
self.mock_run_coroutine.assert_called_once()
async def test_send_reply_no_reply_to(self):
"""Test handling when no reply_to is provided."""
# Execute
await self.handler.send_reply(None, self.mock_data_response)
# Verify
self.mock_run_coroutine.assert_not_called()
async def test_reply_received_callback_with_future(self):
"""Test handling a reply when there's a matching future."""
# Setup
future = asyncio.Future()
self.handler.outstanding["correlation-id"] = future
# Execute
await self.handler.reply_received_callback(self.mock_message)
# Verify
self.assertTrue(future.done())
self.assertEqual(future.result(), self.mock_data_response.body())
self.mock_message.ack.assert_called_once_with(multiple=False)
async def test_reply_received_callback_no_future(self):
"""Test handling a reply when there's no matching future."""
# Execute
await self.handler.reply_received_callback(self.mock_message)
# Verify
self.mock_message.ack.assert_called_once_with(multiple=False)
def test_ensure_trace_info(self):
"""Test ensuring trace info is properly set."""
# Execute
self.handler.ensure_trace_info(self.mock_data_message)
# Verify
self.tracer.start_span.assert_called_once()
def test_record_duration(self):
"""Test recording the duration of message processing."""
# Setup
start_time = 1234567880.0 # 10 seconds before mock_time
# Execute
self.handler.record_duration(start_time, 123)
# No specific assertions needed as this is mostly logging
def test_collect_trace_info(self):
"""Test collecting trace info."""
# Execute
result = collect_trace_info()
# Verify it returns a dictionary
self.assertIsInstance(result, dict)
def test_get_default_trace_state(self):
"""Test getting default trace state."""
# Execute
result = get_default_trace_state()
# Verify it returns a TraceState
self.assertIsNotNone(result)
def test_set_text(self):
"""Test setting text in a carrier."""
# Setup
carrier = {}
# Execute
set_text(carrier, "key", "value")
# Verify
self.assertEqual(carrier["key"], "value")
async def test_reconstitute_data_message_stream_none_result(self):
"""Test reconstituting a message from a stream when the result is None."""
# Setup
self.mock_message.headers = {"clevermicro.addressing": 1}
self.mock_from_stream.return_value = None
# Execute
result = await self.handler.reconstitute_data_message(self.mock_message)
# Verify
self.assertIsNone(result)
self.mock_from_stream.assert_called_once()
self.file_handler.on_message.assert_not_called()
async def test_run_http_request_magic_file_removal_exception(self):
"""Test handling of an exception when removing files."""
# Setup
self.service_adapter.on_message.return_value = self.mock_data_response
multipart_message = MultipartDataMessage(
self.mock_data_message, {"file1": "/tmp/file1.txt"}
)
self.mock_os_remove.side_effect = OSError("File removal error")
# Execute
await self.handler.run_http_request_magic(self.mock_message, multipart_message, self.loop)
# Verify
self.mock_os_remove.assert_called_once_with("/tmp/file1.txt")
# Should continue without raising an exception
async def test_send_reply_content_type_list(self):
"""Test sending a reply with content type as a list."""
# Setup
self.mock_data_response.content_type.return_value = ["application/json", "text/plain"]
# Execute
await self.handler.send_reply("reply-queue", self.mock_data_response)
# Verify
self.mock_run_coroutine.assert_called_once()
# We verify that Message is created with correct content type by inspecting the RabbitMQ call args if available.
self.assertEqual(
"application/json", self.reply_to_exchange.publish.call_args[1]["message"].content_type
)
if __name__ == "__main__":
unittest.main()
+87
View File
@@ -0,0 +1,87 @@
import asyncio
import unittest
from unittest.mock import AsyncMock, MagicMock, patch
from amqp.service.amq_service import AMQService
class TestAMQService(unittest.IsolatedAsyncioTestCase):
def setUp(self):
# Mock the configuration
self.config = MagicMock()
self.config.amq_adapter.service_name = "test-service"
self.config.amq_adapter.swarm_task_id = "test-task"
self.config.amq_adapter.swarm_service_id = "test-service-id"
self.config.backpressure.threshold_load = 100
# Mock the service adapter
self.service_adapter = MagicMock()
# Patch the get_unique_instance_id function
with patch("amqp.service.amq_service.get_unique_instance_id", return_value=12345):
# Create the AMQService instance
self.service = AMQService(self.config, self.service_adapter)
# Mock the backpressure handler with AsyncMock for the coroutine method
self.service.backpressure_handler = MagicMock()
self.service.backpressure_handler.update_backpressure_value = AsyncMock()
self.service.backpressure_handler.loop = asyncio.get_event_loop()
async def test_backpressure_with_maximum(self):
"""Test backpressure method with a specified maximum value."""
# Call the backpressure method with a maximum value
self.service.backpressure(current_availability=50, maximum=200)
# Check that run_coroutine_threadsafe was called with the correct arguments
# We need to patch asyncio.run_coroutine_threadsafe to capture its arguments
with patch("asyncio.run_coroutine_threadsafe") as mock_run:
self.service.backpressure(current_availability=50, maximum=200)
mock_run.assert_called_once()
# The second argument should be the loop
loop_arg = mock_run.call_args[1]["loop"]
self.assertEqual(loop_arg, self.service.backpressure_handler.loop)
# We can also verify the arguments passed to update_backpressure_value
self.service.backpressure_handler.update_backpressure_value.assert_called_with(50, 200)
async def test_backpressure_without_maximum(self):
"""Test backpressure method without a maximum value."""
# Set a maximum_availability value
self.service.maximum_availability = 100
with patch("asyncio.run_coroutine_threadsafe") as mock_run:
# Call the backpressure method without a maximum value
self.service.backpressure(current_availability=50)
# Verify the handler was called with the right values
self.service.backpressure_handler.update_backpressure_value.assert_called_with(50, 100)
mock_run.assert_called_once()
async def test_backpressure_with_negative_maximum(self):
"""Test backpressure method with a negative maximum value."""
# Set a maximum_availability value
self.service.maximum_availability = 100
with patch("asyncio.run_coroutine_threadsafe") as mock_run:
# Call the backpressure method with a negative maximum value
self.service.backpressure(current_availability=50, maximum=-1)
# Verify the handler was called with the right values
self.service.backpressure_handler.update_backpressure_value.assert_called_with(50, 100)
mock_run.assert_called_once()
async def test_backpressure_no_handler(self):
"""Test backpressure method when no handler is available."""
# Set the backpressure handler to None
self.service.backpressure_handler = None
# Call the backpressure method
self.service.backpressure(current_availability=50)
# Nothing should happen, no exception should be raised
if __name__ == "__main__":
unittest.main()
+1 -1
View File
@@ -57,7 +57,7 @@ def main():
parser = argparse.ArgumentParser(description="Create a RabbitMQ exchange.")
parser.add_argument(
"--host",
default="localhost",
default="127.0.0.1",
help="RabbitMQ host (default: localhost)",
)