Compare commits

..

93 Commits

Author SHA1 Message Date
Stanislav Hejny 07f6f4569c feat/#63 - add iterator instead of RabbitMQ callback
Unit test coverage / pytest (pull_request) Failing after 1m49s
Build and Publish Docker Image / build-and-push (push) Successful in 2m29s
Unit test coverage / pytest (push) Failing after 1m45s
/ build-and-push (push) Successful in 1m52s
CI for pypl publish / publish-lib (push) Successful in 1m44s
2025-07-27 09:50:25 +01:00
Stanislav Hejny 66132db13a feat/#63 - cleverswarm unidirectional command
Unit test coverage / pytest (pull_request) Successful in 1m38s
Unit test coverage / pytest (push) Successful in 1m52s
CI for pypl publish / publish-lib (push) Successful in 2m7s
/ build-and-push (push) Successful in 2m20s
Build and Publish Docker Image / build-and-push (push) Successful in 2m32s
2025-07-26 15:41:47 +01:00
Stanislav Hejny 9984501fee feat/#63 - cleverswarm unidirectional command
Unit test coverage / pytest (push) Failing after 1m46s
/ build-and-push (push) Successful in 1m36s
Unit test coverage / pytest (pull_request) Failing after 1m15s
2025-07-26 15:30:12 +01:00
Stanislav Hejny 41fbe46129 feat/#63 - cleverswarm unidirectional command
Unit test coverage / pytest (push) Has been cancelled
/ build-and-push (push) Has been cancelled
Unit test coverage / pytest (pull_request) Failing after 1m51s
2025-07-26 15:29:07 +01:00
Stanislav Hejny 9491517696 feat/#63 - fix the send_command and add unit tests
Unit test coverage / pytest (pull_request) Successful in 1m15s
Unit test coverage / pytest (push) Failing after 1m50s
CI for pypl publish / publish-lib (push) Successful in 1m57s
/ build-and-push (push) Successful in 1m59s
Build and Publish Docker Image / build-and-push (push) Successful in 2m12s
2025-07-24 10:05:36 +01:00
Stanislav Hejny 4985344e5e test: add package_name parameter to invoke_command calls in test_rpc.py
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-24 08:16:42 +01:00
Stanislav Hejny f3be2fbfea test: update tests to use invoke_command directly 2025-07-24 08:16:41 +01:00
Stanislav Hejny ad49cb2836 test: convert demo script to unittest framework in test_rpc.py
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-24 08:08:00 +01:00
Stanislav Hejny 2b20c9a0dc test: add tests for RPC command invocation and error handling 2025-07-24 08:07:57 +01:00
Stanislav Hejny 982c72205a Merge branch 'feat/61-multi-user-support' into develop
Unit test coverage / pytest (push) Successful in 2m16s
/ build-and-push (push) Successful in 2m16s
CI for pypl publish / publish-lib (push) Failing after 2m15s
Build and Publish Docker Image / build-and-push (push) Successful in 2m34s
2025-07-23 19:47:31 +01:00
Stanislav Hejny c3346bbe7a Merge branch 'feat-58-backpressure-reactive-streams' into develop
Unit test coverage / pytest (push) Successful in 1m58s
CI for pypl publish / publish-lib (push) Successful in 2m4s
/ build-and-push (push) Successful in 2m22s
Build and Publish Docker Image / build-and-push (push) Successful in 2m58s
2025-07-23 18:41:51 +01:00
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
hurui200320 80a55f9894 ci: fix docker missing issue
Unit test coverage / pytest (push) Successful in 2m0s
/ build-and-push (push) Successful in 2m4s
Build and Publish Docker Image / build-and-push (push) Successful in 3m5s
Unit test coverage / pytest (pull_request) Successful in 1m7s
base: node:lts-bookworm (debian 12)
2025-07-21 18:29:03 +08:00
Stanislav Hejny d8ab0b861d feat/#58 - fix unit test and deployment pipeline
Unit test coverage / pytest (push) Successful in 1m55s
Unit test coverage / pytest (pull_request) Successful in 1m52s
/ build-and-push (push) Successful in 1m58s
Build and Publish Docker Image / build-and-push (push) Failing after 5m12s
2025-07-21 10:52:46 +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 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 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
87 changed files with 9305 additions and 1434 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
+88
View File
@@ -0,0 +1,88 @@
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 }}
# setup docker
- name: set up docker cli
run: |
apt-get update
apt-get install ca-certificates curl
install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/debian/gpg -o /etc/apt/keyrings/docker.asc
chmod a+r /etc/apt/keyrings/docker.asc
echo \
"deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/debian \
$(. /etc/os-release && echo "$VERSION_CODENAME") stable" | \
tee /etc/apt/sources.list.d/docker.list > /dev/null
apt-get update
apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
- 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
+4 -4
View File
@@ -40,11 +40,11 @@ jobs:
# publish # publish
- name: publish to pulp pypi - name: publish to pulp pypi
run: | run: |
twine upload --repository-url https://pkg.cleverthis.io/pypi/clevermicro/amq-adapter-python/simple/ \ twine upload --repository-url https://git.cleverthis.com/api/packages/clevermicro/pypi \
--username $CI_REGISTRY_USER \ --username ${CI_REGISTRY_USER} \
--password $CI_REGISTRY_PASSWORD \ --password ${CI_REGISTRY_PASSWORD} \
--verbose \ --verbose \
dist/* dist/*
env: env:
CI_REGISTRY_USER: ${{ secrets.REGISTRY_USER }} CI_REGISTRY_USER: ${{ secrets.REGISTRY_USER }}
CI_REGISTRY_PASSWORD: ${{ secrets.REGISTRY_PASSWORD }} CI_REGISTRY_PASSWORD: ${{ secrets.REGISTRY_PASSWORD }}
+2 -1
View File
@@ -56,7 +56,8 @@ amq_adapter.egg-info/
build/ build/
unused-code/ unused-code/
.vscode/
amq_adapter.egg-info/ amq_adapter.egg-info/
build/ build/
coverage.xml 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. 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
```
+167 -177
View File
@@ -3,8 +3,8 @@ import json
import time import time
from asyncio import AbstractEventLoop from asyncio import AbstractEventLoop
from threading import Thread from threading import Thread
from typing import Optional
import psutil
from aio_pika import Message from aio_pika import Message
from aio_pika.abc import AbstractRobustChannel from aio_pika.abc import AbstractRobustChannel
@@ -20,15 +20,15 @@ class ScaleRequestV1:
self, self,
serviceId: str, serviceId: str,
taskId: str, taskId: str,
maxAvailability: int, max_availability: int,
currentAvailability: int, current_availability: int,
requestType: ScalingRequestAlert, requestType: ScalingRequestAlert,
): ):
self.version = 1 # Version of the request, currently 1 self.version = 1 # Version of the request, currently 1
self.serviceId = serviceId self.serviceId = serviceId
self.taskId = taskId self.taskId = taskId
self.maxAvailability = maxAvailability self.max_availability = max_availability
self.currentAvailability = currentAvailability self.current_availability = current_availability
self.requestType = requestType self.requestType = requestType
self.thread = None self.thread = None
@@ -39,56 +39,23 @@ class ScaleRequestV1:
"version": self.version, "version": self.version,
"serviceId": self.serviceId, "serviceId": self.serviceId,
"taskId": self.taskId, "taskId": self.taskId,
"maxAvailability": self.maxAvailability, "maxAvailability": self.max_availability,
"currentAvailability": self.currentAvailability, "currentAvailability": self.current_availability,
"requestType": self.requestType.value, # Using .value for Enum serialization "requestType": self.requestType.value, # Using .value for Enum serialization
}, },
indent=2, 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: class BackpressureHandler:
# Track the number of messages currently being processed # Track the number of messages currently being processed
current_parallel_executions = 0 current_availability = 0
# helps detect IDLE condition # helps detect IDLE condition
last_data_message_time = 0
# helps prevent flooding the system with backpressure events # helps prevent flooding the system with backpressure events
last_backpressure_event_time = 0 last_backpressure_event_time = 0
last_backpressure_event = ScalingRequestAlert.UPDATE last_backpressure_event = ScalingRequestAlert.UPDATE
# _callback_list is used in unit tests to record the invoked callbacks
_callback_list = None
def __init__( def __init__(
self, self,
@@ -103,147 +70,165 @@ class BackpressureHandler:
self.exchange = None self.exchange = None
self.swarm_service_id = self.config.amq_adapter.swarm_service_id self.swarm_service_id = self.config.amq_adapter.swarm_service_id
self.swarm_task_id = self.config.amq_adapter.swarm_task_id self.swarm_task_id = self.config.amq_adapter.swarm_task_id
self.parallel_workers = self.config.backpressure.threshold_threads self.do_loop = -1
self.average_cpu_usage = None 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
self.idle_state_detected_time = 0
def increase_parallel_executions(self): def increase_current_load(self):
"""Increase the number of parallel executions""" """Increase the number of parallel executions, or current_availability load"""
logging_info( logging_info(
"Backpressure: Increase parallel executions, current=%s", "Backpressure: Increase current load (%s) by 1",
self.current_parallel_executions, self.current_load,
) )
self.current_parallel_executions += 1 self.current_load += 1
def decrease_parallel_executions(self): def decrease_current_load(self):
"""Decrease the number of parallel executions""" """Decrease the number of parallel executions, or current load"""
logging_info( logging_info(
"Backpressure: Decrease parallel executions, current=%s", "Backpressure: Decrease current load(%s) by 1",
self.current_parallel_executions, self.current_load,
) )
if self.current_parallel_executions > 0: if self.current_load > 0:
self.current_parallel_executions -= 1 self.current_load -= 1
def update_parallel_executions(self, count: int): def update_current_availability(self, count: int):
"""Update the number of parallel executions""" """Update the number of parallel executions, or current load"""
self.current_parallel_executions = count self.current_availability = count
def update_last_data_message_time(self): def update_last_backpressure_event_time(self):
logging_info("Backpressure: Update last data message time") # logging_info("Backpressure: Update last data message time")
"""Update the 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 # Start the Backpressure monitor loop
self.thread = Thread(target=self.backpressure_monitor_loop) self.thread = Thread(target=self.backpressure_monitor_loop)
self.thread.daemon = True # This makes it a daemon thread
self.thread.start() self.thread.start()
return self.thread return self.thread
async def check_overload_condition(self): 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)
logging_info( or if the service has been idle for too long with high availability (IDLE).
"Backpressure: Check overload condition, current=%s, max=%s",
self.current_parallel_executions, Note: current_availability represents available capacity, not used capacity.
self.parallel_workers, - Low availability (close to 0) means OVERLOAD
) - High availability (close to maximum) with no activity means IDLE
if self.current_parallel_executions >= self.parallel_workers - 1: """
current_time = time.time()
last_event_delta: float = current_time - self.last_backpressure_event_time
# logging_info(
# "Backpressure: Check conditions, current_availability=%s, max=%s, last_activity=%s",
# self.current_availability,
# self.max_availability,
# last_event_delta,
# )
# 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 # Check if the last backpressure event was not an overload
if self.last_backpressure_event != ScalingRequestAlert.OVERLOAD: if (
self.last_backpressure_event != ScalingRequestAlert.OVERLOAD
or last_event_delta > self.config.backpressure.idle_duration
):
# Trigger the overload event # Trigger the overload event
await self.handle_backpressure_overload_event() await self.handle_backpressure_overload_event()
self.last_backpressure_event_time = time.time() # 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 self.last_backpressure_event = ScalingRequestAlert.OVERLOAD
def backpressure_monitor_loop(self): # Check for IDLE condition - high availability (more than 80% of maximum) with no activity
_time_window_millis = self.config.backpressure.time_window elif self.current_availability >= round(0.8 * self.max_availability):
_idle_duration_millis = self.config.backpressure.idle_duration idle_duration = self.config.backpressure.idle_duration
_overload_duration_millis = self.config.backpressure.idle_duration # Check if service has been idle for longer than the configured duration
_loop = asyncio.new_event_loop() if (
_monitor_interval = 0.2 # Monitor every 200ms current_time - self.idle_state_detected_time > idle_duration
CPU_IDLE, CPU_ACTIVE, CPU_OVERLOAD = 0, 1, 2 and self.last_backpressure_event == ScalingRequestAlert.IDLE
IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 15, 90 ):
self.average_cpu_usage = RunningAverage( logging_info(
int( "Backpressure: Service has been idle for %s seconds (threshold: %s seconds)",
max(_overload_duration_millis, _idle_duration_millis) last_event_delta,
// _monitor_interval idle_duration,
)
# Trigger the idle event
await self._handle_backpressure_idle_event()
# update / reset time-window so that the IDLE is not sent too often
self.last_backpressure_event_time = current_time
self.idle_state_detected_time = current_time
# send IDLE even only once
self.last_backpressure_event = ScalingRequestAlert.UPDATE
else:
# Don't send IDLE right away when the availability increases, but wait for a while, send UPDATE instead
self.last_backpressure_event = ScalingRequestAlert.IDLE
if current_time - self.idle_state_detected_time > idle_duration:
self.idle_state_detected_time = current_time
if 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
# 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
async def _backpressure_monitor(): async def _backpressure_monitor(self):
"""Monitor the backpressure conditions and trigger events accordingly""" """Periodically monitor the backpressure conditions and trigger events accordingly"""
_cpu_usage_state = CPU_ACTIVE _monitor_interval = 0.5 # Monitor every 500ms
_cpu_usage_changed = time.time() while self._do_loop():
while True: await self.check_overload_condition()
_current_cpu_usage = psutil.cpu_percent(interval=None) await asyncio.sleep(_monitor_interval)
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 def backpressure_monitor_loop(self):
if ( _loop = asyncio.new_event_loop()
_cpu_usage_state == CPU_OVERLOAD _loop.run_until_complete(self._backpressure_monitor())
and _now - _cpu_usage_changed > _overload_duration_millis
and (
_now - self.last_backpressure_event_time > _time_window_millis
or self.last_backpressure_event != ScalingRequestAlert.OVERLOAD
)
):
# Trigger the overload event
await self.handle_backpressure_overload_event()
self.last_backpressure_event_time = _now
self.last_backpressure_event = ScalingRequestAlert.OVERLOAD
elif (
_cpu_usage_state == CPU_IDLE
and _now - _cpu_usage_changed > _idle_duration_millis
and _now - self.last_data_message_time > _idle_duration_millis
and (
_now - self.last_backpressure_event_time > _time_window_millis
or self.last_backpressure_event != ScalingRequestAlert.IDLE
)
):
# Trigger the idle event
await self._handle_backpressure_idle_event()
self.last_backpressure_event = ScalingRequestAlert.IDLE
self.last_backpressure_event_time = _now
else:
# Trigger update event
if _now - self.last_backpressure_event_time > _time_window_millis:
# Trigger the update event
_scaling_request: ScaleRequestV1 = ScaleRequestV1(
self.swarm_service_id,
self.swarm_task_id,
100,
_average_cpu_usage,
ScalingRequestAlert.UPDATE,
)
# Address the message to any adapter capable of supporting BACKPRESSURE request
await self.publish_backpressure_request(_scaling_request)
self.last_backpressure_event_time = _now
self.last_backpressure_event = ScalingRequestAlert.UPDATE
await asyncio.sleep(_monitor_interval)
_loop.run_until_complete(_backpressure_monitor())
async def handle_backpressure_overload_event(self): async def handle_backpressure_overload_event(self):
logging_warning("Backpressure: Capacity close to depleted!") logging_warning("Backpressure: Capacity close to depleted!")
@@ -251,14 +236,12 @@ class BackpressureHandler:
scaling_request: ScaleRequestV1 = ScaleRequestV1( scaling_request: ScaleRequestV1 = ScaleRequestV1(
self.swarm_service_id, self.swarm_service_id,
self.swarm_task_id, self.swarm_task_id,
self.parallel_workers, self.max_availability,
self.parallel_workers - self.current_parallel_executions, self.current_availability, # Current availability is passed directly
ScalingRequestAlert.OVERLOAD, ScalingRequestAlert.OVERLOAD,
) )
# Address the message to any management service instance capable of supporting BACKPRESSURE request # Address the message to any management service instance capable of supporting BACKPRESSURE request
await self.publish_backpressure_request(scaling_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): async def _handle_backpressure_idle_event(self):
logging_warning("Backpressure: Service is idle.") logging_warning("Backpressure: Service is idle.")
@@ -266,14 +249,12 @@ class BackpressureHandler:
scaling_request: ScaleRequestV1 = ScaleRequestV1( scaling_request: ScaleRequestV1 = ScaleRequestV1(
self.swarm_service_id, self.swarm_service_id,
self.swarm_task_id, self.swarm_task_id,
self.parallel_workers, self.max_availability,
self.parallel_workers - self.current_parallel_executions, self.current_availability, # Current availability is passed directly
ScalingRequestAlert.IDLE, ScalingRequestAlert.IDLE,
) )
# Address the message to any adapter capable of supporting BACKPRESSURE request # Address the message to any adapter capable of supporting BACKPRESSURE request
await self.publish_backpressure_request(scaling_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): async def publish_backpressure_request(self, scaling_request: ScaleRequestV1):
# Publish the backpressure request to the management service # Publish the backpressure request to the management service
@@ -283,15 +264,13 @@ class BackpressureHandler:
if not self.exchange: if not self.exchange:
async def _wrap_rabbit_mq_api_init(channel): async def _wrap_rabbit_mq_api_init(channel):
_exchange = await channel.get_exchange( _exchange = await channel.get_exchange(name="cleverthis.clevermicro.management")
name="cleverthis.clevermicro.management"
)
return _exchange return _exchange
if BackpressureHandler._callback_list is not None:
BackpressureHandler._callback_list["_wrap_rabbit_mq_api_init"] = 1
self.exchange = await await_result( self.exchange = await await_result(
asyncio.run_coroutine_threadsafe( asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api_init(self.channel), self.loop)
_wrap_rabbit_mq_api_init(self.channel), self.loop
)
) )
if self.exchange: if self.exchange:
@@ -319,6 +298,17 @@ class BackpressureHandler:
return True return True
return False return False
await await_future( if BackpressureHandler._callback_list is not None:
asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api(), self.loop) 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
+517 -180
View File
@@ -4,55 +4,71 @@ Methods here will invoke the actual CleverThis Service code, which are passed as
""" """
import asyncio import asyncio
import enum
import inspect import inspect
import json import json
import pathlib
import re
from asyncio import AbstractEventLoop 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 ( from aio_pika.abc import (
AbstractExchange,
AbstractRobustChannel, AbstractRobustChannel,
AbstractRobustConnection, AbstractRobustConnection,
AbstractExchange,
) )
from aiohttp import ClientSession, ClientResponse from aiohttp import ClientResponse, ClientSession
from dataclasses import dataclass from fastapi import HTTPException, UploadFile
from fastapi import HTTPException from fastapi.params import File
from fastapi.routing import APIRoute
from fastapi.security import OAuth2PasswordRequestForm
from pydantic import BaseModel from pydantic import BaseModel
from starlette.datastructures import Headers
from starlette.responses import FileResponse from starlette.responses import FileResponse
from amqp.adapter import serializer, pydantic_serializer from amqp.adapter import serializer
from amqp.adapter.data_parser import AMQDataParser from amqp.adapter.data_parser import AMQDataParser
from amqp.adapter.data_response_factory import DataResponseFactory from amqp.adapter.data_response_factory import DataResponseFactory
from amqp.adapter.file_handler import FileHandler, StreamingFileHandler 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.config.amq_configuration import AMQConfiguration
from amqp.model.model import DataMessage, DataResponse, AMQErrorMessage from amqp.model.model import AMQErrorMessage, AMQMessage, DataMessage, DataResponse
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE
from amqp.router.utils import await_result 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 Create objects for reply.
in case the response should initiate file download. Return (channel for reply queue, exchange for reply, reply queue name)
""" """
def __init__(self, data_message: DataMessage, connection: AbstractRobustConnection): # Wrap calls to RabbitMQ API to make it Awaitable, so it can be run thread-safe with
# Copy all fields from DataMessage # correct event loop. This is needed because the RabbitMQ API is not thread-safe and needs
self.magic = data_message.magic # to be called in the context of event loop the first RabbitMQ connection was created with.
self.id = data_message.id _channel = await connection.channel()
self.method = data_message.method _exchange_name: str = AMQ_REPLY_TO_EXCHANGE
self.domain = data_message.domain _exchange = await _channel.get_exchange(name=_exchange_name, ensure=True)
self.path = data_message.path _queue = await _channel.declare_queue(exclusive=False) # Declare a unique, non-exclusive queue
self.headers = data_message.headers _queue_name = _queue.name
self.trace_info = data_message.trace_info await _queue.bind(
self.base64_body = data_message.base64_body exchange=_exchange_name,
# Add the new field routing_key=_queue_name, # Use queue name as routing key
self.connection = connection )
self.extra_data = ( logging_debug(f"Publishing to exchange: {_exchange_name}, q: {_queue_name}")
data_message.extra_data if hasattr(data_message, "extra_data") else {} return _channel, _exchange, _queue_name
)
async def _convert_file_response( async def _convert_file_response(
@@ -63,41 +79,22 @@ async def _convert_file_response(
) -> str: ) -> str:
""" """
Converts a FileResponse object to a JSON string containing the filename and stream name. Converts a FileResponse object to a JSON string containing the filename and stream name.
The file content itself is transmitted through the named stream (which is implemented as temporary dedicated queue. The file content itself is transmitted through the named stream (which is implemented as
temporary dedicated queue.
""" """
# All calls to RabbitMQ API have to be called threadsafe, because those appear to be attached
async def _wrap_amq_api_calls( # to parent event loop, and need to execute at that event loop, not the current_availability one
connection: AbstractRobustConnection, # allow coroutine to execute. calling _future.result() will block the entire thread
) -> tuple[AbstractRobustChannel, AbstractExchange, str]: # and the coroutine will never execute
# Wrap calls to RabbitMQ API to make it Awaitable, so it can be run thread-safe with correct event loop
# This is needed because the RabbitMQ API is not thread-safe and needs to be called in the context of
# the event loop the first RabbitMQ connection was created with.
_channel = await connection.channel()
_exchange_name: str = AMQ_REPLY_TO_EXCHANGE
_exchange = await _channel.get_exchange(name=_exchange_name, ensure=True)
_queue = await _channel.declare_queue(
exclusive=False
) # Declare a unique, non-exclusive queue
_queue_name = _queue.name
await _queue.bind(
exchange=_exchange_name,
routing_key=_queue_name, # Use queue name as routing key
)
logging_debug(f"Publishing to exchange: {_exchange_name}, q: {_queue_name}")
return _channel, _exchange, _queue_name
# All calls to RabbitMQ API have to be called threadsafe, because those appear to be attached to parent
# event loop, and need to execute at that event loop, not the current one
# allow coroutine to execute. calling _future.result() will block the entire thread and the coroutine will never execute
_channel, _exchange, _routing_key = await await_result( _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 # publish_file also needs RabbitMQ API, ensure the loop is passed as argument
(_stream_name, _file_size) = await file_handler.publish_file( (_stream_name, _file_size) = await file_handler.publish_file(
_exchange, obj.path, _routing_key, loop _exchange, obj.path, _routing_key, loop
) )
# as above, but luckily this time we don't need to wait for the result # 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( return json.dumps(
{ {
"files": [ "files": [
@@ -114,6 +111,8 @@ async def _convert_file_response(
def is_likely_json(text: str) -> bool: def is_likely_json(text: str) -> bool:
"""Quick test to see if should convert the response string to JSON format.""" """Quick test to see if should convert the response string to JSON format."""
if not isinstance(text, str):
return False
_text = text.strip() _text = text.strip()
return (_text.startswith("{") and _text.endswith("}")) or ( return (_text.startswith("{") and _text.endswith("}")) or (
_text.startswith("[") and _text.endswith("]") _text.startswith("[") and _text.endswith("]")
@@ -123,20 +122,28 @@ def is_likely_json(text: str) -> bool:
# ============= CleverThisServiceAdapter ============= # ============= CleverThisServiceAdapter =============
class CleverThisServiceAdapter: class CleverThisServiceAdapter:
""" """
A base class for CleverThis service adapter that handles the incoming AMQP messages and routes them to the A base class for CleverThis service adapter that handles the incoming AMQP messages and routes
appropriate CleverThisService API endpoint. 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 IMPORTANT: This class is not intended to be used directly. It should be subclassed for
and override the methods to handle the specific API calls. each CleverThis service and override the methods to handle the specific API calls.
""" """
def __init__( 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. Initializes the CleverThisServiceAdapter with the given AMQ configuration and optional HTTP
The HTTP session is for the case when the adapter uses loose coupling with the CleverThis service via session. The HTTP session is for the case when the adapter uses loose coupling with
HTTP REST API. 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.amq_configuration = amq_configuration
self.session = session self.session = session
self.service_host = self.amq_configuration.dispatch.service_host self.service_host = self.amq_configuration.dispatch.service_host
@@ -145,6 +152,81 @@ class CleverThisServiceAdapter:
self.require_authenticated_user = ( self.require_authenticated_user = (
self.amq_configuration.amq_adapter.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: def get_content_type(self, message_headers: Dict[str, List[str]]) -> str:
""" """
@@ -159,9 +241,11 @@ class CleverThisServiceAdapter:
async def get_current_user(self, token: str) -> object: 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, To Be Overriden in subclass for given CleverThis Service, to return the current_availability user based
in format appropriate for the service. on the token, in format appropriate for the service.
:param token: The token to be used for authentication. :param token: The token to be used for authentication.
:return: A user object or None if the user is not authenticated. :return: A user object or None if the user is not authenticated.
""" """
return None return None
@@ -175,11 +259,9 @@ class CleverThisServiceAdapter:
_auth_user: Any | None = None _auth_user: Any | None = None
_amq_response: Any | None = None _amq_response: Any | None = None
logging_debug( logging_debug(f"on_message: require_authenticated_user={self.require_authenticated_user}")
f"on_message: require_authenticated_user={self.require_authenticated_user}"
)
if 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}") logging_debug(f"Auth: {_auth}")
if isinstance(_auth, List): if isinstance(_auth, List):
_auth = _auth[0] _auth = _auth[0]
@@ -191,12 +273,8 @@ class CleverThisServiceAdapter:
except Exception as error: except Exception as error:
logging_error(f"on_message: Error getting user: {error}") logging_error(f"on_message: Error getting user: {error}")
if ( if _auth_user or not self.require_authenticated_user or message.path().endswith("/login"):
_auth_user method = message.method().lower()
or not self.require_authenticated_user
or message.path.endswith("/login")
):
method = message.method.lower()
if method == "get": if method == "get":
_amq_response = await self.get(message, _auth_user) _amq_response = await self.get(message, _auth_user)
elif method == "put": elif method == "put":
@@ -208,162 +286,422 @@ class CleverThisServiceAdapter:
elif method == "delete": elif method == "delete":
_amq_response = await self.delete(message, _auth_user) _amq_response = await self.delete(message, _auth_user)
else: else:
raise ValueError(f"Unexpected HTTP method: {message.method}") raise ValueError(f"Unexpected HTTP method: {message.method()}")
else: else:
_amq_response = DataResponseFactory.of_error_message( _amq_response = DataResponseFactory.of_error_message(
message.id, message.id(),
AMQErrorMessage(401, "Unauthorized", "Unauthorized"), AMQErrorMessage(401, "Unauthorized", "Unauthorized"),
message.trace_info, message.trace_info(),
) )
logging_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 return _amq_response
# ================================================================================================== async def _authorize_forward_request(
# ===================== C l e v e r T h i s A P I m e t h o d s ============================ self, message: DataMessage, request_coroutine, auth_user: Any
# ================================================================================================== ) -> DataResponse:
async def get(self, message: DataMessage, auth_user: Any) -> DataResponse:
""" """
Handles the GET requests for the CleverThis API. to be overriden in subclass. This is a base implementation. This method is used to authorize the request and forward it to the CleverThis Service API.
It will attempt to call the REST endpoint directly using HTTP protocol. 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 message: message from the AMQP that contains the GET request
:param auth_user: user authenticated by the API Gateway, as UserSchem :param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP :return: DataResponse to be sent back to the API Gateway via AMQP
""" """
# unmatched path - try to invoke the service directly on this path if self.session is None:
url = f"http://{self.service_host}:{self.service_port}{message.path}" return await self.handle_no_body_payload(message, auth_user, "GET")
headers = {**message.headers, **message.trace_info}
# 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( logging_info(
f"SEND HEADERS out: {', '.join(f'{k}={v}' for k, v in headers.items())}, REQ={url}" 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)
_http_resp: ClientResponse = await self.session.get(url, headers=headers) return DataResponseFactory.of(
return DataResponseFactory.of( message.id(),
message.id, _http_resp.status,
_http_resp.status, _http_resp.content_type,
_http_resp.content_type, await _http_resp.read(),
await _http_resp.read(), message.trace_info(),
message.trace_info,
)
return DataResponseFactory.of_error_message(
message.id,
AMQErrorMessage(404, "Not Found", "Destination location does not exists"),
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. Handles the PUT requests for the CleverThis API. Uses path, path params, query string and
It will attempt to call the REST endpoint directly using HTTP protocol. 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 message: message from the AMQP that contains the PUT request
:param auth_user: user authenticated by the API Gateway, as UserSchem :param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP :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, message,
self.session.put( self.session.put(
f"http://{self.service_host}:{self.service_port}{message.path}", f"http://{self.service_host}:{self.service_port}{message.path()}",
data=message.body, data=message.body(),
headers={**message.headers, **message.trace_info}, headers={**message.headers(), **message.trace_info()},
), ),
auth_user, 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. Handles the POST requests for the CleverThis API. Uses path, path params, query string and
It will attempt to call the REST endpoint directly using HTTP protocol. body payload to determine the endpoint method to invoke and its input arguments
:param message: message from the AMQP that contains the POST request, including body and headers 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 :param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP :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, message,
self.session.post( self.session.post(
f"http://{self.service_host}:{self.service_port}{message.path}", f"http://{self.service_host}:{self.service_port}{message.path()}",
data=message.body, data=message.body(),
headers={**message.headers, **message.trace_info}, headers={**message.headers(), **message.trace_info()},
), ),
auth_user, auth_user,
) )
async def handle_possible_form( async def head(self, message: AMQMessage, auth_user: Any) -> DataResponse:
self, message: DataMessage, request_coroutine, auth_user: Any """
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: ) -> 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. Handles the requests that do not have a body payload.
It will attempt to call the REST endpoint directly using HTTP protocol. :param message: message from the AMQP that contains the request
: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 :param auth_user: user authenticated by the API Gateway, as UserSchem
:return: DataResponse to be sent back to the API Gateway via AMQP :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. Handles the requests that have a body payload, including Multipart data.
It will attempt to call the REST endpoint directly using HTTP protocol. :param message: message from the AMQP that contains the request
: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 :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 :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 _verified_args = {}
async def call_cleverthis_api( _required_args = getattr(_route.endpoint, "__annotations__", {})
message: AMQMessage, _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, args: Any,
api_method: Callable[[Any, Any], Coroutine], api_method: Callable[[Any, Any], Coroutine],
success_code: int, ) -> dict:
loop: AbstractEventLoop | None = None,
) -> DataResponse:
""" """
Handles POST and PUT requests, which typically contain possibly complex payload data, passed as a dictionary. Check the provided arguments against the API method signature and convert the incoming string value
Invokes the provided API method with the given arguments and authenticated user, into correct expected input object.
and returns the response as DataResponse. The logic used FastAPI annotation and would use the value from 'default_factory' invocation (where defined)
:param message: DataMessage containing the details of the call and the payload or 'default' value, as defined in the FastAPI annotation, to create default value in case the object is not in
:param args: map of the arguments to pass to the API method the input values provided by the caller.
:param api_method: Callable - the method to invoke. Last parameter is the authenticated user
:param success_code: HTTP code to be returned when operation suceeds
:return: DataResponse class
""" """
_verified_args = {}
try: try:
_verified_args = {}
_sig = inspect.signature(api_method) _sig = inspect.signature(api_method)
_required_args = list(_sig.parameters.values()) _required_args = list(_sig.parameters.values())
# _required_args = getattr(api_method, '__annotations__', {}) # _required_args = getattr(api_method, '__annotations__', {})
for arg in _required_args: for arg in _required_args:
if arg.name in args: if arg.name in args:
_annotation = arg.annotation
_str_annotation = str(_annotation)
# convert to another type if arg is not a str # convert to another type if arg is not a str
if arg.annotation == int: if arg.name == self.authenticated_user_arg_name:
_verified_args[arg.name] = int(args[arg.name]) _verified_args[arg.name] = args.get(arg.name)
elif arg.annotation == float: elif _annotation == int:
_verified_args[arg.name] = float(args[arg.name]) _verified_args[arg.name] = int(args.get(arg.name))
elif arg.annotation == bool: elif _annotation == float:
_verified_args[arg.name] = serializer.str_to_bool( _verified_args[arg.name] = float(args.get(arg.name))
args[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
) )
# TODO: more? 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: 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: else:
logging_info( if hasattr(arg, "default"):
f"{api_method.__name__}: Missing required argument: {arg.name}" _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( 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) _result: Any = await api_method(**_verified_args)
logging_info( logging_info(f"ENDPOINT RESULT: [{api_method.__name__}], RESULT: {str(_result)}")
f"ENDPOINT RESULT: {api_method.__name__}, RESULT: {str(_result)}"
)
if isinstance(_result, FileResponse): if isinstance(_result, FileResponse):
_json_result, _content_type = ( _json_result, _content_type = (
@@ -371,7 +709,7 @@ class CleverThisServiceAdapter:
_result, _result,
message.connection, message.connection,
file_handler=StreamingFileHandler(), file_handler=StreamingFileHandler(),
loop=loop, loop=self.loop,
), ),
"application/octet-stream", "application/octet-stream",
) )
@@ -379,7 +717,10 @@ class CleverThisServiceAdapter:
_json_result = _result.json() _json_result = _result.json()
_content_type = "application/json" _content_type = "application/json"
else: else:
_json_result = _result if isinstance(_result, str):
_json_result = _result
else:
_json_result = python_type_to_json(_result)
_content_type = "application/json" _content_type = "application/json"
_binary_result = ( _binary_result = (
@@ -388,60 +729,56 @@ class CleverThisServiceAdapter:
else ( else (
_json_result.body _json_result.body
if hasattr(_json_result, "body") if hasattr(_json_result, "body")
else ( else (_json_result.read() if hasattr(_json_result, "read") else _json_result)
_json_result.read()
if hasattr(_json_result, "read")
else _json_result
)
) )
) )
logging_info( 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( return DataResponseFactory.of(
message.id, message.id(),
success_code, success_code,
_content_type, _content_type,
_binary_result, _binary_result,
message.trace_info, message.trace_info(),
) )
except HTTPException as http_error: except HTTPException as http_error:
logging_error(f"Service REST endpoint invocation failed: {http_error}") logging_error(f"Service REST endpoint invocation failed: {http_error}")
_content = str(http_error.detail) _content = str(http_error.detail)
if not is_likely_json(_content): if not is_likely_json(_content):
return DataResponseFactory.of_error_message( return DataResponseFactory.of_error_message(
message.id, message.id(),
AMQErrorMessage( AMQErrorMessage(
http_error.status_code, http_error.status_code,
"Service Invocation Failed", "Service Invocation Failed",
_content, _content,
), ),
message.trace_info, message.trace_info(),
) )
return DataResponseFactory.of( return DataResponseFactory.of(
message.id, message.id(),
http_error.status_code, http_error.status_code,
"application/json", "application/json",
_content.encode("utf-8"), _content.encode("utf-8"),
message.trace_info, message.trace_info(),
) )
except Exception as error: except Exception as error:
logging_error(f"Service endpoint invocation failed: {error}") logging_error(f"Service endpoint invocation failed: {error}")
_content = str(error) _content = str(error)
if not is_likely_json(_content): if not is_likely_json(_content):
return DataResponseFactory.of_error_message( return DataResponseFactory.of_error_message(
message.id, message.id(),
AMQErrorMessage( AMQErrorMessage(
500, 500,
"Service Invocation Failed", "Service Invocation Failed",
_content, _content,
), ),
message.trace_info, message.trace_info(),
) )
return DataResponseFactory.of( return DataResponseFactory.of(
message.id, message.id(),
500, 500,
"application/json", "application/json",
_content.encode("utf-8"), _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()
+86 -40
View File
@@ -5,21 +5,22 @@ methods to create, serialize and deserialize DataMessage objects.
import base64 import base64
import json import json
import re import re
from io import BytesIO
from datetime import datetime, timezone 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 import amqp.adapter.file_handler
from amqp.model.model import DataMessage from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
from amqp.model.snowflake_id import SnowflakeId from amqp.adapter.pydantic_serializer import python_type_to_json
from amqp.adapter.serializer import ( from amqp.adapter.serializer import (
map_with_list_as_string,
map_as_string, map_as_string,
map_with_list_as_string,
parse_map_list_string, parse_map_list_string,
parse_map_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.router.utils import sanitize_as_rabbitmq_name
from amqp.adapter.file_handler import StreamingFileHandler, FileHandler
class DataMessageFactory: class DataMessageFactory:
@@ -30,7 +31,6 @@ class DataMessageFactory:
MACHINE_ID_MASK = (1 << MACHINE_ID_BITS) - 1 MACHINE_ID_MASK = (1 << MACHINE_ID_BITS) - 1
SEQUENCE_MASK = (1 << SEQUENCE_BITS) - 1 SEQUENCE_MASK = (1 << SEQUENCE_BITS) - 1
MAGIC_BYTE_HTTP_REQUEST = "A"
_instance = None _instance = None
snowflake_epoch = datetime(2024, 10, 1, 0, 0, tzinfo=timezone.utc).timestamp() snowflake_epoch = datetime(2024, 10, 1, 0, 0, tzinfo=timezone.utc).timestamp()
@@ -100,12 +100,10 @@ class DataMessageFactory:
:return: DataMessage object :return: DataMessage object
""" """
serializable_headers = headers.copy() if headers else {} serializable_headers = headers.copy() if headers else {}
payload = base64.b64encode( payload = base64.b64encode(message.encode("utf-8") if isinstance(message, str) else message)
message.encode("utf-8") if isinstance(message, str) else message
)
trace_info = {} trace_info = {}
return DataMessage( return DataMessage(
self.MAGIC_BYTE_HTTP_REQUEST, DataMessageMagicByte.HTTP_REQUEST.value,
self.generate_snowflake_id(), self.generate_snowflake_id(),
method, method,
domain, domain,
@@ -115,6 +113,56 @@ class DataMessageFactory:
payload, payload,
) )
def safe_rsplit(self, fq_method_name, default="_"):
"""
Safely split a fully qualified method name into package, class, and method
Args:
fq_method_name (str): Fully qualified method name
default (str, optional): Default value for missing tokens. Defaults to ''.
Returns:
tuple: (package, class_name, method)
"""
if not fq_method_name:
return default, default, default
parts = fq_method_name.rsplit(".", 2)
# Pad with default values if insufficient tokens
while len(parts) < 3:
parts.insert(0, default)
return tuple(parts)
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 = self.safe_rsplit(fq_method_name)
headers = {"Content-Type": ["application/octet-stream"], "Return-Address": [swarm_id]}
trace_info = {}
body = python_type_to_json(args).encode(encoding="utf-8") if args else b""
print(
f"Creating RPC message for method: {fq_method_name}, args: {args} body: {body.decode('utf-8') if body else 'None'}"
)
return DataMessage(
DataMessageMagicByte.RPC_REQUEST.value,
self.generate_snowflake_id(),
package,
class_name,
method,
headers,
trace_info,
body,
)
@staticmethod @staticmethod
def amq_message_routing_key(message: DataMessage) -> str: def amq_message_routing_key(message: DataMessage) -> str:
""" """
@@ -122,7 +170,7 @@ class DataMessageFactory:
:param message: DataMessage input :param message: DataMessage input
:return: routing key (as string) compliant to RabbitMQ allowed character set. :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 @staticmethod
def get_timestamp_from_id(_id: SnowflakeId) -> int: def get_timestamp_from_id(_id: SnowflakeId) -> int:
@@ -132,9 +180,7 @@ class DataMessageFactory:
:return: the message timestamp :return: the message timestamp
""" """
bits = _id.get_bits() bits = _id.get_bits()
low = bits[0] >> ( low = bits[0] >> (DataMessageFactory.SEQUENCE_BITS + DataMessageFactory.MACHINE_ID_BITS)
DataMessageFactory.SEQUENCE_BITS + DataMessageFactory.MACHINE_ID_BITS
)
high = bits[1] << ( high = bits[1] << (
64 - DataMessageFactory.SEQUENCE_BITS - DataMessageFactory.MACHINE_ID_BITS 64 - DataMessageFactory.SEQUENCE_BITS - DataMessageFactory.MACHINE_ID_BITS
) )
@@ -149,13 +195,13 @@ class DataMessageFactory:
""" """
return ( return (
f"DataMessage{{" f"DataMessage{{"
f"id={message.id}, " f"id={message.id()}, "
f"method='{message.method}', " f"method='{message.method()}', "
f"domain='{message.domain}', " f"domain='{message.domain()}', "
f"path='{message.path}', " f"path='{message.path()}', "
f"headers={{{map_with_list_as_string(message.headers)}}}, " f"headers={{{map_with_list_as_string(message.headers())}}}, "
f"traceInfo={{{map_as_string(message.trace_info)}}}, " f"traceInfo={{{map_as_string(message.trace_info())}}}, "
f"base64body={message.base64_body.decode()}" f"base64body={message.base64body().decode()}"
f"}}" f"}}"
) )
@@ -166,20 +212,20 @@ class DataMessageFactory:
:param message: message to serialize :param message: message to serialize
:return: messages as bytes :return: messages as bytes
""" """
id_bytes = str(message.id).encode() id_bytes = str(message.id()).encode()
prolog = ( 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") ).encode("utf-8")
header_length = 1 + len(id_bytes) + len(prolog) header_length = 1 + len(id_bytes) + len(prolog)
prolog_len = header_length.to_bytes(2, byteorder="big") prolog_len = header_length.to_bytes(2, byteorder="big")
msg_length = 2 + header_length + len(message.base64_body) msg_length = 2 + header_length + len(message.base64body())
with BytesIO(bytearray(msg_length)) as baos: baos = BytesIO(bytearray(msg_length))
baos.write(prolog_len) baos.write(prolog_len)
baos.write(DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST.encode("utf-8")) baos.write(message.magic().encode("utf-8"))
baos.write(id_bytes) baos.write(id_bytes)
baos.write(prolog) baos.write(prolog)
baos.write(message.base64_body) baos.write(message.base64body())
return baos.getvalue() return baos.getvalue()
@staticmethod @staticmethod
def from_bytes(input_bytes: bytes) -> DataMessage: def from_bytes(input_bytes: bytes) -> DataMessage:
@@ -192,7 +238,7 @@ class DataMessageFactory:
metadata = input_bytes[2 : prolog_len + 2].decode() metadata = input_bytes[2 : prolog_len + 2].decode()
pattern = re.compile( pattern = re.compile(
f"{DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST}" r"([AR])"
r"([0-9A-Fa-f]+)\|" r"([0-9A-Fa-f]+)\|"
r"m=([^|]*)\|" r"m=([^|]*)\|"
r"d=([^|]*)\|" r"d=([^|]*)\|"
@@ -208,19 +254,19 @@ class DataMessageFactory:
f"Input string [{metadata}] does not match pattern: {pattern.pattern}" f"Input string [{metadata}] does not match pattern: {pattern.pattern}"
) )
id_str = match.group(1) id_str = match.group(2)
method = match.group(2) method = match.group(3)
domain = match.group(3) domain = match.group(4)
path = match.group(4) path = match.group(5)
headers_str = match.group(5) headers_str = match.group(6)
trace_info_str = match.group(6) trace_info_str = match.group(7)
body_b64 = input_bytes[prolog_len + 2 :] body_b64 = input_bytes[prolog_len + 2 :]
headers = parse_map_list_string(headers_str) headers = parse_map_list_string(headers_str)
trace_info = parse_map_string(trace_info_str) trace_info = parse_map_string(trace_info_str)
return DataMessage( return DataMessage(
DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST, match.group(1), # Magic byte (R for RPC, A for HTTP)
SnowflakeId.from_hex(id_str), SnowflakeId.from_hex(id_str),
method, method,
domain, domain,
+21 -33
View File
@@ -1,16 +1,15 @@
import io import io
import json import json
import python_multipart from typing import Dict, Tuple
from tempfile import SpooledTemporaryFile
from fastapi import UploadFile
from python_multipart.multipart import Field, File
from typing import Tuple, Dict
from urllib.parse import parse_qs, urlparse from urllib.parse import parse_qs, urlparse
from xml.etree import ElementTree from xml.etree import ElementTree
from amqp.adapter.logging_utils import logging_error, logging_debug import python_multipart
from amqp.model.model import DataMessage 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: class MultipartFormDataParser:
@@ -23,15 +22,14 @@ class MultipartFormDataParser:
included, the message is transmitted in the unchanged multipart/form-data format. 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.form_data: dict = {}
self.is_multipart = False self.is_multipart = False
self.is_json = False self.is_json = False
self.is_valid = True self.is_valid = True
# Convert each list of strings to a single string joined by newline, then to bytes # Convert each list of strings to a single string joined by newline, then to bytes
_headers = { _headers = {
key: "\n".join(value).encode("utf-8") key: "\n".join(value).encode("utf-8") for key, value in message.headers().items()
for key, value in message.headers.items()
} }
try: try:
if ( if (
@@ -42,7 +40,7 @@ class MultipartFormDataParser:
try: try:
self.form_data = json.loads(message.body()) self.form_data = json.loads(message.body())
self.is_json = True 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) self.form_data = AMQDataParser.process_form_data(self.form_data)
except Exception as jsonerror: except Exception as jsonerror:
logging_error(f"Parsing JSON: {jsonerror}") logging_error(f"Parsing JSON: {jsonerror}")
@@ -53,15 +51,14 @@ class MultipartFormDataParser:
) )
self.is_multipart = True self.is_multipart = True
except Exception as e: except Exception as e:
logging_error( __error = f"Error parsing multipart/form-data: {e}. Trying parsing as JSON."
f"Error parsing multipart/form-data: {e}. Trying parsing as JSON."
)
try: try:
self.form_data = json.loads(message.body()) self.form_data = json.loads(message.body())
self.is_json = True 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) self.form_data = AMQDataParser.process_form_data(self.form_data)
except Exception as jsonerror: except Exception as jsonerror:
logging_error(__error)
logging_error(f"Parsing JSON: {jsonerror}") logging_error(f"Parsing JSON: {jsonerror}")
self.is_valid = False self.is_valid = False
@@ -72,11 +69,8 @@ class MultipartFormDataParser:
def on_file(self, file: File): def on_file(self, file: File):
logging_debug(str(file)) logging_debug(str(file))
file.file_object.seek(0) file.file_object.seek(0)
_spooled_file: SpooledTemporaryFile = SpooledTemporaryFile()
_spooled_file.write(file.file_object.read())
_spooled_file.seek(0)
self.form_data[file.field_name.decode("utf-8")] = UploadFile( self.form_data[file.field_name.decode("utf-8")] = UploadFile(
_spooled_file, size=file.size, filename=file.file_name.decode("utf-8") file.file_object, size=file.size, filename=file.file_name.decode("utf-8")
) )
@@ -116,7 +110,6 @@ class AMQDataParser:
for key, value in data["files"].items(): for key, value in data["files"].items():
if value: # Only override if files data is non-empty if value: # Only override if files data is non-empty
result[key] = value[0] result[key] = value[0]
return result return result
@staticmethod @staticmethod
@@ -135,13 +128,11 @@ class AMQDataParser:
_parsed = urlparse(url) _parsed = urlparse(url)
_query_params = parse_qs(_parsed.query) _query_params = parse_qs(_parsed.query)
# Convert values from lists to single values when there's only one value # Convert values from lists to single values when there's only one value
_simplified_params = { _simplified_params = {k: v[0] if len(v) == 1 else v for k, v in _query_params.items()}
k: v[0] if len(v) == 1 else v for k, v in _query_params.items()
}
return _parsed.path, _simplified_params return _parsed.path, _simplified_params
@staticmethod @staticmethod
def parse_request_body(message: DataMessage): def parse_request_body(message: AMQMessage):
""" """
Parses the HTTP POST request body based on the MIME type. 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. This is typically called fom the concrete CleverThis Service to extract parameters from the message body.
@@ -151,14 +142,14 @@ class AMQDataParser:
Returns: Returns:
dict: Parsed data as a dictionary. dict: Parsed data as a dictionary.
""" """
_content_type = message.headers.get("Content-Type", "") _content_type = message.headers().get("Content-Type", "")
if not _content_type: if not _content_type:
raise ValueError("Content-Type header is required") raise ValueError("Content-Type header is required")
content_type = _content_type[0] content_type = _content_type[0]
# Parse JSON # Parse JSON
if content_type == "application/json": if content_type == "application/json":
try: try:
return json.loads(message.body()) return json.loads(message.body() or "{}")
except json.JSONDecodeError as e: except json.JSONDecodeError as e:
logging_error(f"Invalid JSON: {e}") logging_error(f"Invalid JSON: {e}")
raise ValueError(f"Invalid JSON: {e}") raise ValueError(f"Invalid JSON: {e}")
@@ -167,15 +158,12 @@ class AMQDataParser:
elif content_type == "application/x-www-form-urlencoded": elif content_type == "application/x-www-form-urlencoded":
try: try:
_form_data: dict = { _form_data: dict = {
k: v[0] if len(v) == 1 else v k: v[0] if len(v) == 1 else v for k, v in parse_qs(message.body()).items()
for k, v in parse_qs(message.body()).items()
} }
if len(_form_data) == 0 and len(message.body()) > 0: if len(_form_data) == 0 and len(message.body()) > 0:
_form_data = json.loads(message.body()) _form_data = json.loads(message.body())
if _form_data["files"] and _form_data["form"]: if _form_data["files"] is not None and _form_data["form"] is not None:
_form_data = MultipartFormDataParser.process_form_data( _form_data = AMQDataParser.process_form_data(_form_data)
_form_data
)
return _form_data return _form_data
except Exception as e: except Exception as e:
logging_error(f"Invalid URL-encoded form data: {e}") logging_error(f"Invalid URL-encoded form data: {e}")
+29 -21
View File
@@ -5,13 +5,13 @@ method to create, serialize and deserialize DataResponse message
import base64 import base64
import re import re
from io import BytesIO from io import BytesIO
from amqp.model.model import DataResponse, AMQErrorMessage
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.model.snowflake_id import SnowflakeId
from amqp.adapter.serializer import parse_map_string, map_as_string
class DataResponseFactory: class DataResponseFactory:
MAGIC_BYTE_HTTP_REQUEST = "A"
@staticmethod @staticmethod
def create_async_response_message(_id: SnowflakeId) -> DataResponse: def create_async_response_message(_id: SnowflakeId) -> DataResponse:
@@ -20,8 +20,14 @@ class DataResponseFactory:
:param _id: ID :param _id: ID
:return: Response message :return: Response message
""" """
return DataResponse( return DataResponseFactory.of(
str(_id), 200, "application/text", "OK".encode("utf-8"), "", "", {} id=_id,
response_code=200,
content_type="application/text",
body="OK".encode("utf-8"),
trace_info={},
error="",
error_cause="",
) )
@staticmethod @staticmethod
@@ -44,13 +50,15 @@ class DataResponseFactory:
:return: DataResponseMessage object :return: DataResponseMessage object
""" """
return DataResponse( return DataResponse(
id=id.as_string(), magic=DataMessageMagicByte.HTTP_REQUEST.value,
response_code=response_code, id=id,
response_code=str(response_code),
content_type=content_type, content_type=content_type,
response=body,
error=error, error=error,
error_cause=error_cause, error_cause=error_cause,
headers={},
trace_info=trace_info, trace_info=trace_info,
base64body=base64.b64encode(body),
) )
@staticmethod @staticmethod
@@ -60,11 +68,11 @@ class DataResponseFactory:
trace_info: dict[str, str], trace_info: dict[str, str],
) -> DataResponse: ) -> DataResponse:
return DataResponseFactory.of( return DataResponseFactory.of(
id, id=id,
error_message.response_code, response_code=error_message.response_code,
error_message.content_type, content_type=error_message.content_type,
error_message.serialize(), body=error_message.serialize(),
trace_info, trace_info=trace_info,
error=error_message.error, error=error_message.error,
error_cause=error_message.detail, error_cause=error_message.detail,
) )
@@ -76,19 +84,19 @@ class DataResponseFactory:
:param response: Object to serialize :param response: Object to serialize
:return: byte array as serialized message :return: byte array as serialized message
""" """
id_bytes = response.id.encode("utf-8") id_bytes = response.id().as_string().encode("utf-8")
prolog = ( 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") ).encode("utf-8")
header_length = 1 + len(id_bytes) + len(prolog) header_length = 1 + len(id_bytes) + len(prolog)
prolog_len = header_length.to_bytes(2, byteorder="big") prolog_len = header_length.to_bytes(2, byteorder="big")
msg_length = 2 + header_length + len(response.response) msg_length = 2 + header_length + len(response.base64body())
with BytesIO(bytearray(msg_length)) as baos: with BytesIO(bytearray(msg_length)) as baos:
baos.write(prolog_len) 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(id_bytes)
baos.write(prolog) baos.write(prolog)
baos.write(base64.b64encode(response.response)) baos.write(response.base64body())
return baos.getvalue() return baos.getvalue()
@staticmethod @staticmethod
@@ -102,7 +110,7 @@ class DataResponseFactory:
metadata = input_bytes[2 : prolog_len + 2].decode() metadata = input_bytes[2 : prolog_len + 2].decode()
pattern = re.compile( pattern = re.compile(
f"{DataResponseFactory.MAGIC_BYTE_HTTP_REQUEST}" f"{DataMessageMagicByte.HTTP_REQUEST.value}"
r"([0-9A-Fa-f]+)\|" r"([0-9A-Fa-f]+)\|"
r"r=([^|]*)\|" r"r=([^|]*)\|"
r"c=([^|]*)\|" r"c=([^|]*)\|"
@@ -128,14 +136,14 @@ class DataResponseFactory:
trace_info = parse_map_string(trace_info_str) trace_info = parse_map_string(trace_info_str)
return DataResponse( return DataResponseFactory.of(
id_str, id_str,
response_code, response_code,
content_type, content_type,
base64.b64decode(body_b64), base64.b64decode(body_b64),
trace_info,
error, error,
error_cause, error_cause,
trace_info,
) )
+86 -80
View File
@@ -4,22 +4,22 @@ implements the file downloader for the amqp adapter, where file is sent in chunk
import asyncio import asyncio
import json import json
import os
import logging import logging
from asyncio import Future, AbstractEventLoop import os
from asyncio import AbstractEventLoop
from collections import defaultdict from collections import defaultdict
from pathlib import Path from pathlib import Path
from aio_pika import connect_robust, Message from aio_pika import Message, connect_robust
from aio_pika.abc import ( from aio_pika.abc import (
AbstractIncomingMessage,
AbstractRobustChannel,
AbstractQueue,
AbstractExchange, AbstractExchange,
AbstractIncomingMessage,
AbstractQueue,
AbstractRobustChannel,
DeliveryMode, 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 from amqp.router.utils import await_result
# Global dictionary to track completed downloads # Global dictionary to track completed downloads
@@ -153,6 +153,44 @@ class StreamingFileHandler(FileHandler):
def __init__(self): def __init__(self):
super().__init__() 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( async def download_file(
self, loop, rabbitmq_url, file_info, message_data, output_dir self, loop, rabbitmq_url, file_info, message_data, output_dir
) -> (int, int): ) -> (int, int):
@@ -166,7 +204,6 @@ class StreamingFileHandler(FileHandler):
message_data (dict): The entire message message_data (dict): The entire message
output_dir (str): The directory where the file should be saved. output_dir (str): The directory where the file should be saved.
""" """
global download_locations
filename = file_info["filename"] filename = file_info["filename"]
size = file_info["size"] size = file_info["size"]
stream_name = file_info["streamName"] stream_name = file_info["streamName"]
@@ -185,41 +222,46 @@ class StreamingFileHandler(FileHandler):
f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name}:{queue_name} to {filepath}" 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: if queue_name:
_local_file_size, _local_eof = await await_result( _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 _local_file_size, _local_eof
return 0, CANCELLED return 0, CANCELLED
async def download_buffer(self, loop, rabbitmq_url, queue_name) -> (bytearray, int): async def _download_from_queue_to_buffer(
self, rabbitmq_url, queue_name, loop
) -> tuple[bytearray, int]:
"""
Internal method for downloading a file from a queue to RAM.
Return (bytearray, eof status).
"""
_connection = await self._get_connection(rabbitmq_url, loop=loop)
_eof = IN_PROGRESS
_res = bytearray()
try:
channel: AbstractRobustChannel = await _connection.channel()
queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
logging_debug(f"Awaiting all buffer chunks being downloaded for {queue_name}")
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
_res.extend(message.body)
_eof = message.headers.get("eof", END_OF_FILE)
if _eof != IN_PROGRESS:
break
logging_info(f"Finished downloading buffer {queue_name}")
except Exception as e:
logging_error(f"Downloading buffer: {e}")
await _connection.close()
return _res, _eof
async def download_buffer(self, loop, rabbitmq_url, queue_name: str | None) -> (bytearray, int):
""" """
Asynchronously downloads a bytearay from RabbitMQ based on the provided file information. Asynchronously downloads a bytearay from RabbitMQ based on the provided file information.
It assumes entire array is in the first (and only) chunk It assumes entire array is in the first (and only) chunk
@@ -230,35 +272,10 @@ class StreamingFileHandler(FileHandler):
""" """
logging_debug(f"Downloading buffer from stream: {queue_name}") logging_debug(f"Downloading buffer from stream: {queue_name}")
if queue_name: if queue_name:
async def wrap_rabbit_mq_calls() -> tuple[bytearray, int]:
_connection = await connect_robust(rabbitmq_url, loop=loop)
_eof = IN_PROGRESS
_res = bytearray()
try:
channel: AbstractRobustChannel = await _connection.channel()
queue: AbstractQueue = await channel.get_queue(
queue_name, ensure=False
)
logging_debug(
f"Awaiting all buffer chunks being downloaded for {queue_name}"
)
async with queue.iterator() as queue_iter:
async for message in queue_iter:
async with message.process():
_res.extend(message.body)
_eof = message.headers.get("eof", END_OF_FILE)
if _eof != IN_PROGRESS:
break
logging_info(f"Finished downloading buffer {queue_name}")
except Exception as e:
logging_error(f"Downloading buffer: {e}")
await _connection.close()
return _res, _eof
_local_res, _local_eof = await await_result( _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 _local_res, _local_eof
return 0, CANCELLED return 0, CANCELLED
@@ -281,7 +298,6 @@ class StreamingFileHandler(FileHandler):
loop: The asyncio event loop of the RabbitMQ API loop: The asyncio event loop of the RabbitMQ API
output_dir output_dir
""" """
global download_locations
try: try:
_message_data = json.loads(message.body.decode()) _message_data = json.loads(message.body.decode())
_amq_message_data = json.loads(json_data.decode()) _amq_message_data = json.loads(json_data.decode())
@@ -298,22 +314,16 @@ class StreamingFileHandler(FileHandler):
files_to_download.append(files_info) files_to_download.append(files_info)
if not files_to_download: if not files_to_download:
logging_info( logging_info(" [%s] No files to download in this message.", message.delivery_tag)
" [%s] No files to download in this message.", message.delivery_tag
)
# await message.ack() # await message.ack()
return defaultdict() return defaultdict()
# Create a task for each file download # Create a task for each file download
_tasks = [] _tasks = []
for file_info in files_to_download: for file_info in files_to_download:
logging_info( logging_info(f" [{message.delivery_tag}] about to create task for {file_info}")
f" [{message.delivery_tag}] about to create task for {file_info}"
)
_task = asyncio.create_task( _task = asyncio.create_task(
self.download_file( self.download_file(loop, rabbitmq_url, file_info, _message_data, output_dir)
loop, rabbitmq_url, file_info, _message_data, output_dir
)
) )
_tasks.append(_task) _tasks.append(_task)
logging_info( logging_info(
@@ -379,9 +389,7 @@ class StreamingFileHandler(FileHandler):
"total_size": _file_size, "total_size": _file_size,
"chunk_size": len(_chunk), "chunk_size": len(_chunk),
"eof": ( "eof": (
END_OF_FILE END_OF_FILE if max_chunk_size + _offset >= _file_size else IN_PROGRESS
if max_chunk_size + _offset >= _file_size
else IN_PROGRESS
), ),
}, },
) )
@@ -394,9 +402,7 @@ class StreamingFileHandler(FileHandler):
while not _future.done(): while not _future.done():
# sleep allows coroutine to execute. must wait for completion to ensure all chunks written in right order # sleep allows coroutine to execute. must wait for completion to ensure all chunks written in right order
await asyncio.sleep(0.02) await asyncio.sleep(0.02)
logging_debug( logging_debug(f"Waiting for chunk {_offset}/{_file_size} future= {_future}")
f"Waiting for chunk {_offset}/{_file_size} future= {_future}"
)
_offset += len(_chunk) _offset += len(_chunk)
logging_debug( logging_debug(
+8 -8
View File
@@ -13,17 +13,17 @@ __verbose__ = os.environ.get("AMQ_VERBOSE_LOGGING", "0") == "1"
def get_context_info(): def get_context_info():
""" """
Retrieves the current module, line number, and function name. Retrieves the current_availability module, line number, and function name.
Returns: Returns:
A tuple containing (module name, line number, function name). A tuple containing (module name, line number, function name).
Returns (None, None, None) if it fails to retrieve the information. Returns (None, None, None) if it fails to retrieve the information.
""" """
try: try:
# Get the frame object for the current function call # Get the frame object for the current_availability function call
frame = inspect.currentframe() frame = inspect.currentframe()
if frame is not None: 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 frame = frame.f_back.f_back
if frame: if frame:
code_context = frame.f_code code_context = frame.f_code
@@ -44,7 +44,7 @@ def get_context_info():
def logging_error(msg: str, *args) -> None: 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. Add module name, line and function name where the error occurred, on single line.
Args: Args:
msg (str): The error message to log. msg (str): The error message to log.
@@ -70,7 +70,7 @@ def logging_error(msg: str, *args) -> None:
def logging_warning(msg: str, *args) -> None: 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. Add module name, line and function name where the print occurred, on single line.
Args: Args:
msg (str): The warning message to log. msg (str): The warning message to log.
@@ -91,7 +91,7 @@ def logging_warning(msg: str, *args) -> None:
def logging_info(msg: str, *args) -> None: 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. Add module name, line and function name where the print occurred, on single line.
Args: Args:
msg (str): The info message to log. msg (str): The info message to log.
@@ -112,7 +112,7 @@ def logging_info(msg: str, *args) -> None:
def logging_debug(msg: str, *args) -> None: 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. Add module name, line and function name where the print occurred, on single line.
Args: Args:
msg (str): The debug message to log. msg (str): The debug message to log.
@@ -120,7 +120,7 @@ def logging_debug(msg: str, *args) -> None:
_thread_id = threading.get_ident() _thread_id = threading.get_ident()
module, line_number, function_name = get_context_info() module, line_number, function_name = get_context_info()
if module and line_number and function_name: if module and line_number and function_name:
logging.info( logging.debug(
f"{_thread_id}:{__version__} [DBG] '{msg}' ---> '{module}':{line_number}, '{function_name}()'", f"{_thread_id}:{__version__} [DBG] '{msg}' ---> '{module}':{line_number}, '{function_name}()'",
*args, *args,
) )
+93 -83
View File
@@ -2,12 +2,15 @@
This module provides serialization and deserialization functions for Pydantic models that extend BaseModel This module provides serialization and deserialization functions for Pydantic models that extend BaseModel
""" """
import uuid import importlib
import json import json
from datetime import datetime import uuid
from typing import Dict, Any 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 pydantic import BaseModel
from urllib.parse import urlparse, parse_qs
from amqp.adapter.logging_utils import logging_debug from amqp.adapter.logging_utils import logging_debug
@@ -25,8 +28,12 @@ def serialize_object(obj: BaseModel) -> str:
A string representation of the object. A string representation of the object.
""" """
class_name = obj.__class__.__name__ class_name = obj.__class__.__name__
module_name = obj.__class__.__module__
obj_dict = {} 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): if isinstance(value, uuid.UUID):
obj_dict[key] = str(value) obj_dict[key] = str(value)
elif isinstance(value, datetime): elif isinstance(value, datetime):
@@ -46,8 +53,50 @@ def serialize_object(obj: BaseModel) -> str:
else: else:
obj_dict[key] = value obj_dict[key] = value
dict_string = ",".join(f"{k}:{v}" for k, v in obj_dict.items()) dict_string = json.dumps(obj_dict)
return f"{class_name}:{{{dict_string}}}" 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: def deserialize_object(data: str) -> BaseModel:
@@ -62,93 +111,54 @@ def deserialize_object(data: str) -> BaseModel:
Returns: Returns:
An R2RSerializable object. An R2RSerializable object.
""" """
class_name, dict_str = data.split(":", 1) class_path, dict_str = data.split(":", 1)
dict_str = dict_str.strip("{}") module_path, class_name = class_path.rsplit(".", 1)
pairs = dict_str.split(",") if dict_str else []
obj_dict: Dict[str, Any] = {} obj_dict: Dict[str, Any] = json.loads(dict_str)
for pair in pairs:
if not pair or ":" not in pair:
continue
k, v = pair.split(":", 1)
v = v.strip() # important
obj_dict[k] = _convert_value(v)
logging_debug(f"Deserialized object: {class_name}") logging_debug(f"Deserialized object: {class_path}")
if "[" in class_name: if "[" in class_name:
class_name = class_name.split("[")[0] class_name = class_name.split("[")[0]
logging_debug(f"trying object: {class_name}") 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: if cls is None:
raise ValueError(f"Unknown class: {class_name}") raise ValueError(f"Unknown class: {class_name}")
# process nested model, dict and list
for key, value in obj_dict.items():
key_type: type[Any] = cls.model_fields[key].annotation
# handle union with none, this will give the true type
# for example, return `A` if the type is `A | None`
if get_origin(key_type) is UnionType or get_origin(key_type) is Union:
union_types = get_args(key_type)
if len(union_types) > 2:
# if there are multiple unions, like `A | B | None`
# then we have no way to figure out which one to use
raise ValueError(f"Unsupported Union type: {key_type}")
# return the none-None one as the type for processing
if union_types[0] is type(None):
key_type = union_types[1]
else:
key_type = union_types[0]
# skip if the value is None
if value is None:
continue
if key_type == uuid.UUID:
obj_dict[key] = uuid.UUID(value)
elif key_type == datetime:
obj_dict[key] = datetime.fromisoformat(value)
elif key_type == dict:
obj_dict[key] = value
elif issubclass(key_type, BaseModel):
obj_dict[key] = deserialize_object(value)
elif key_type == list:
obj_dict[key] = list(map(try_deserialize_object, value))
return cls(**obj_dict) return cls(**obj_dict)
def _convert_value(v: str) -> Any:
"""
Helper function to convert a string to its appropriate Python type.
"""
if v == "None":
return None
elif _is_uuid_string(v):
return uuid.UUID(v)
elif _is_datetime_string(v):
return datetime.fromisoformat(v)
elif _is_int_string(v):
return int(v)
elif (
v.startswith("{")
and v.endswith("}")
and ":" in v
and "{" not in v[1:-1]
and "}" not in v[1:-1]
): # nested object
# check if it is a json dict
try:
json.loads(v)
return json.loads(v)
except json.JSONDecodeError:
return deserialize_object(v) # recursive call
elif v.startswith("[") and v.endswith("]"): # list
v = v[1:-1] # remove brackets
items = v.split(",") if v else []
deserialized_list = []
for item in items:
item = item.strip()
deserialized_list.append(_convert_value(item)) # recursive call
return deserialized_list
else:
return v # Keep as string
def _is_uuid_string(s: str) -> bool:
"""Helper to check if a string is a valid UUID."""
try:
uuid.UUID(s)
return True
except ValueError:
return False
def _is_datetime_string(s: str) -> bool:
"""Helper to check if a string is a valid datetime in ISO format"""
try:
datetime.fromisoformat(s)
return True
except ValueError:
return False
def _is_int_string(s: str) -> bool:
"""Helper to check if a string is a valid int."""
try:
int(s)
return True
except ValueError:
return False
def parse_url_query(url): def parse_url_query(url):
""" """
Parses a HTTP GET URL and returns (context_path, query_params_dict) Parses a HTTP GET URL and returns (context_path, query_params_dict)
+13 -6
View File
@@ -4,7 +4,7 @@ Helper functions used to aid with (de)serialization of messages sent over AMQP.
import struct import struct
from io import BytesIO from io import BytesIO
from typing import List, Union, Dict from typing import Dict, List, Union
def long_to_bytes(x: int) -> bytes: def long_to_bytes(x: int) -> bytes:
@@ -18,7 +18,7 @@ def bytes_to_long(b: bytes) -> int:
def read_long(bis: BytesIO) -> int: def read_long(bis: BytesIO) -> int:
try: try:
return struct.unpack("q", bis.read(8))[0] return struct.unpack("q", bis.read(8))[0]
except IOError: except struct.error:
return 0 return 0
@@ -39,10 +39,14 @@ def map_with_list_as_string(map_: Dict[str, List[str]]) -> str:
def parse_map_string(input_str: str) -> Dict[str, str]: def parse_map_string(input_str: str) -> Dict[str, str]:
map_ = {} map_ = {}
if input_str.startswith("{") and input_str.endswith("}"): if input_str.startswith("{") and input_str.endswith("}"):
for entry in input_str[1:-1].split(","): for entry in str(input_str[1:-1]).replace("\n", "").split(","):
if len(entry) > 1: if len(entry) > 1:
key, value = entry.split("=") if entry.startswith('"'):
map_[key] = value key, value = entry.split(":")
map_[key.strip('"')] = value.strip('"')
else:
key, value = entry.split("=")
map_[key] = value
return map_ return map_
@@ -51,7 +55,10 @@ def csv_as_list(input_str: str) -> List[str]:
def add_to_map(map_: Dict[str, List[str]], token: str) -> None: def add_to_map(map_: Dict[str, List[str]], token: str) -> None:
eq_idx = token.index("=") try:
eq_idx = token.index("=")
except ValueError:
eq_idx = -1 # set to -1 when not found
if eq_idx > 0: if eq_idx > 0:
key = token[:eq_idx] key = token[:eq_idx]
map_[key] = csv_as_list(token[eq_idx + 1 :]) map_[key] = csv_as_list(token[eq_idx + 1 :])
+4 -13
View File
@@ -1,13 +1,12 @@
import base64 import base64
import logging
from amqp.adapter.data_message_factory import DataMessageFactory from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.logging_utils import logging_error, logging_info 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.adapter.trace_info_adapter import TraceInfoAdapter
from amqp.model.model import ServiceMessage from amqp.model.model import ServiceMessage
from amqp.model.service_message_type import ServiceMessageType from amqp.model.service_message_type import ServiceMessageType
from amqp.model.snowflake_id import SnowflakeId from amqp.model.snowflake_id import SnowflakeId
from amqp.adapter.serializer import parse_map_string, map_as_string
RS = "|" RS = "|"
SPLIT_RS = r"\|" SPLIT_RS = r"\|"
@@ -99,9 +98,7 @@ class ServiceMessageFactory:
i += 1 i += 1
reply_to = input_str_array[i] if i < len(input_str_array) else "" reply_to = input_str_array[i] if i < len(input_str_array) else ""
i += 1 i += 1
trace_info = ( trace_info = parse_map_string(input_str_array[i]) if i < len(input_str_array) else {}
parse_map_string(input_str_array[i]) if i < len(input_str_array) else {}
)
i += 1 i += 1
message = input_str_array[i] if i < len(input_str_array) else "" message = input_str_array[i] if i < len(input_str_array) else ""
return ServiceMessage( return ServiceMessage(
@@ -111,11 +108,7 @@ class ServiceMessageFactory:
recipient_name=recipient_name, recipient_name=recipient_name,
reply_to=reply_to, reply_to=reply_to,
trace_info=trace_info, trace_info=trace_info,
message=( message=(base64.b64decode(message).decode() if payload_type & 0x80 else message),
base64.b64decode(message).decode()
if payload_type & 0x80
else message
),
) )
except Exception as e: except Exception as e:
logging_error( logging_error(
@@ -148,9 +141,7 @@ class ServiceMessageFactory:
:return: formatted message type :return: formatted message type
""" """
num_type = ( num_type = (
message.message_type.value message.message_type.value if message.message_type else ServiceMessageType.UNKNOWN.value
if message.message_type
else ServiceMessageType.UNKNOWN.value
) - 1 ) - 1
fmt = (num_type & 0x7F) | (message.payload_type & 0x80) fmt = (num_type & 0x7F) | (message.payload_type & 0x80)
return f"{fmt:02X}" 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
+91 -18
View File
@@ -5,7 +5,7 @@ import os
from amqp.adapter.logging_utils import logging_error, logging_warning from amqp.adapter.logging_utils import logging_error, logging_warning
""" """
Convert configuration from properties file to an object. Convert configuration from properties file to an object.
""" """
@@ -24,9 +24,7 @@ class AMQAdapter:
:param config: config parser to read configuration values :param config: config parser to read configuration values
""" """
self.generator_id = config.getint( self.generator_id = 0
"CleverMicro-AMQ", self.PREFIX + "generator-id", fallback=164
)
self.service_name = config.get( self.service_name = config.get(
"CleverMicro-AMQ", "CleverMicro-AMQ",
self.PREFIX + "service-name", self.PREFIX + "service-name",
@@ -40,11 +38,77 @@ class AMQAdapter:
self.PREFIX + "require-authenticated-user", self.PREFIX + "require-authenticated-user",
fallback=False, fallback=False,
) )
# these 2 shall be provided by Docker swarm, format is a random alphanumeric string self.backpressure_enabled = config.getboolean(
self.swarm_service_id = os.getenv( "CleverMicro-AMQ",
"SWARM_SERVICE_ID", default="clever-amqp-service" self.PREFIX + "default-backpressure-enabled",
fallback=False,
) )
self.callbacks_enabled = config.getboolean(
"CleverMicro-AMQ",
self.PREFIX + "callbacks-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") 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: def adapter_prefix(self) -> str:
return f"{self.service_name}-{self.generator_id}-" return f"{self.service_name}-{self.generator_id}-"
@@ -68,15 +132,9 @@ class Dispatch:
:param config: config parser to read configuration values :param config: config parser to read configuration values
""" """
self.use_dlq = config.getboolean( self.use_dlq = config.getboolean("CleverMicro-AMQ", self.PREFIX + "use-dlq", fallback=False)
"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_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( self.amq_port_tls = config.getint(
"CleverMicro-AMQ", self.PREFIX + "amq-port-tls", fallback=5671 "CleverMicro-AMQ", self.PREFIX + "amq-port-tls", fallback=5671
) )
@@ -110,7 +168,7 @@ class Backpressure:
:param config: config parser to read configuration values :param config: config parser to read configuration values
""" """
self.threshold_threads = config.getint( self.threshold_load = config.getint(
"CleverMicro-AMQ", "CleverMicro-AMQ",
self.PREFIX + "threshold", self.PREFIX + "threshold",
fallback=int(os.getenv("PARALLEL_WORKERS", 1)), fallback=int(os.getenv("PARALLEL_WORKERS", 1)),
@@ -132,6 +190,10 @@ class Backpressure:
self.idle_duration = config.getint( self.idle_duration = config.getint(
"CleverMicro-AMQ", self.PREFIX + "idle-duration", fallback=30 "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: class AMQConfiguration:
@@ -146,7 +208,7 @@ class AMQConfiguration:
config.read(config_file_name) config.read(config_file_name)
else: else:
# Get the directory of the caller (the module where this function is called) # 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]) caller_module = inspect.getmodule(caller_frame[0])
if caller_module: if caller_module:
caller_dir = os.path.dirname(os.path.abspath(caller_module.__file__)) caller_dir = os.path.dirname(os.path.abspath(caller_module.__file__))
@@ -171,3 +233,14 @@ class AMQConfiguration:
self.amq_adapter: AMQAdapter = AMQAdapter(config) self.amq_adapter: AMQAdapter = AMQAdapter(config)
self.dispatch: Dispatch = Dispatch(config) self.dispatch: Dispatch = Dispatch(config)
self.backpressure: Backpressure = Backpressure(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
+11 -2
View File
@@ -14,6 +14,13 @@ cm.amq-adapter.generator-id=1
cm.amq-adapter.route-mapping= cm.amq-adapter.route-mapping=
# Indicate if AMQ Adapter should respond with HTTP 401 Unauthorized if the request does not contain valid JWT token # 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 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
cm.amq-adapter.callbacks-enabled=False
# CleverMicro AMQ Adapter dispatcher, settings related to RabbitMQ connection and optional REST forwarding # CleverMicro AMQ Adapter dispatcher, settings related to RabbitMQ connection and optional REST forwarding
cm.dispatch.use-dlq=true cm.dispatch.use-dlq=true
@@ -27,7 +34,8 @@ cm.dispatch.download-dir=/tmp/downloaded_files
#cm.dispatch.service-port=8080 #cm.dispatch.service-port=8080
# Backpressure monitor settings # 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 cm.backpressure.threshold=5
# Define maximum CPU usage (%) before triggering backpressure OVERLOAD alert # Define maximum CPU usage (%) before triggering backpressure OVERLOAD alert
cm.backpressure.threshold-cpu-overload=90 cm.backpressure.threshold-cpu-overload=90
@@ -39,4 +47,5 @@ cm.backpressure.time-window=10
cm.backpressure.cpu-overload-duration=30 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 # Define the time window length, in seconds, during which, if no request is made to the Service, the system will trigger IDLE alert
cm.backpressure.cpu-idle-duration=30 cm.backpressure.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)
+97 -69
View File
@@ -1,8 +1,10 @@
import base64 import base64
import json import json
from enum import Enum
from typing import List, Dict
from dataclasses import dataclass, field 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.service_message_type import ServiceMessageType
from amqp.model.snowflake_id import SnowflakeId 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: class CleverMicroMessage:
DEFAULT_CONTENT_TYPE = ["application/octet-stream"] 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__( def __init__(
self, self,
magic: str, magic: str,
@@ -66,11 +88,22 @@ class CleverMicroMessage:
def body(self) -> bytes: def body(self) -> bytes:
return base64.b64decode(self.base64body()) 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] 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): 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__( def __init__(
self, self,
magic: str, magic: str,
@@ -82,9 +115,7 @@ class DataMessage(CleverMicroMessage):
trace_info: Dict[str, str], trace_info: Dict[str, str],
base64body: bytes, base64body: bytes,
): ):
super().__init__( super().__init__(magic, id, method, domain, path, headers, trace_info, base64body)
magic, id, method, domain, path, headers, trace_info, base64body
)
def method(self) -> str: def method(self) -> str:
return self.m_field() return self.m_field()
@@ -95,19 +126,77 @@ class DataMessage(CleverMicroMessage):
def path(self) -> str: def path(self) -> str:
return self.p_field() 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): 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__( def __init__(
self, self,
magic: str, magic: str,
id: SnowflakeId, id: SnowflakeId,
response_code: str, response_code: str,
content_type: str,
error: str, error: str,
error_cause: str, error_cause: str,
headers: Dict[str, List[str]], headers: Dict[str, List[str]],
trace_info: Dict[str, str], trace_info: Dict[str, str],
base64body: bytes, base64body: bytes,
): ):
if headers is None:
headers = {}
headers["Content-Type"] = [content_type]
super().__init__( super().__init__(
magic, magic,
id, id,
@@ -179,65 +268,6 @@ class AMQRoute:
return self.__str__() 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 # Ensure backward compatibility
AMQResponse = DataResponse AMQResponse = DataResponse
@@ -266,7 +296,7 @@ class ScalingRequestAlert(Enum):
IDLE = "IDLE" IDLE = "IDLE"
UPDATE = "UPDATE" UPDATE = "UPDATE"
SCALE_UP = "SCALE_UP" SCALE_UP = "SCALE_UP"
SCSLE_DOWN = "SCALE_DOWN" SCALE_DOWN = "SCALE_DOWN"
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -301,9 +331,7 @@ class ScalingRequest:
task_id=data["taskId"], task_id=data["taskId"],
max_availability=data["maxAvailability"], max_availability=data["maxAvailability"],
current_availability=data["currentAvailability"], current_availability=data["currentAvailability"],
request_type=ScalingRequestAlert( request_type=ScalingRequestAlert(data["requestType"]), # Convert string to Enum
data["requestType"]
), # Convert string to Enum
) )
+36 -37
View File
@@ -1,23 +1,23 @@
import logging import logging
from typing import Optional
import aio_pika import aio_pika
from aio_pika.abc import AbstractRobustQueue, AbstractRobustExchange from aio_pika.abc import AbstractRobustExchange, AbstractRobustQueue
from pika.exchange_type import ExchangeType from pika.exchange_type import ExchangeType
from amqp.adapter.data_message_factory import DataMessageFactory from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.logging_utils import ( from amqp.adapter.logging_utils import (
logging_error,
logging_debug, logging_debug,
logging_error,
logging_info, logging_info,
logging_warning, logging_warning,
) )
from amqp.model.model import DataMessage, AMQRoute from amqp.model.model import AMQRoute, DataMessage
from amqp.router.router_base import ( from amqp.router.router_base import (
AMQ_REPLY_TO_EXCHANGE, AMQ_REPLY_TO_EXCHANGE,
DLQ_EXCHANGE,
AMQ_RPC_EXCHANGE, AMQ_RPC_EXCHANGE,
DLQ_EXCHANGE,
) )
from amqp.router.router_consumer import RouterConsumer
from amqp.router.utils import sanitize_as_rabbitmq_name from amqp.router.utils import sanitize_as_rabbitmq_name
@@ -25,7 +25,7 @@ class RabbitMQClient:
PREFETCH_COUNT = 1 PREFETCH_COUNT = 1
DLQ_ARGS = {"x-dead-letter-exchange": DLQ_EXCHANGE} 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.channel: aio_pika.RobustChannel | None = None
self.connection: aio_pika.RobustConnection | None = None self.connection: aio_pika.RobustConnection | None = None
self.rpc_exchange: aio_pika.RobustExchange | None = None self.rpc_exchange: aio_pika.RobustExchange | None = None
@@ -37,7 +37,7 @@ class RabbitMQClient:
self.amq_configuration.dispatch.amq_host, self.amq_configuration.dispatch.amq_host,
) )
logging.info("*******************************************") logging.info("*******************************************")
self.router = RouterConsumer(configuration) self.router = router
# this.meterRegistry = new PrometheusMeterRegistry(PrometheusConfig.DEFAULT); # this.meterRegistry = new PrometheusMeterRegistry(PrometheusConfig.DEFAULT);
self.reply_to_exchange: AbstractRobustExchange | None = ( self.reply_to_exchange: AbstractRobustExchange | None = (
@@ -85,40 +85,36 @@ class RabbitMQClient:
After this the Service is listening on (consuming from) the queue for incoming messages. After this the Service is listening on (consuming from) the queue for incoming messages.
""" """
if self.router.is_open(self.channel): if self.router.is_open(self.channel):
queue_args = ( queue_args = self.DLQ_ARGS if self.amq_configuration.dispatch.use_dlq else None
self.DLQ_ARGS if self.amq_configuration.dispatch.use_dlq else None logging_info("RabbitMQClient register routes, cfg mapping: [%s]", route_mapping)
)
logging_info(
"RabbitMQClient register routes, cfg mapping: [%s]", route_mapping
)
# convert the comma-separated string into list of AMQRoute objects # convert the comma-separated string into list of AMQRoute objects
requested_routes = self.router.unwrap_route_list(route_mapping) requested_routes = self.router.unwrap_route_list(route_mapping)
for route in requested_routes: for route in requested_routes:
try: try:
# Note: AMQRoute allows one of the exchange or queue to be None, in which case the default is used # Note: AMQRoute allows one of the exchange or queue to be None, in which case the default is used
queue_name = ( queue_name = (
route.queue route.queue if route.queue else self.router.get_consume_queue_name()
if route.queue
else self.router.get_consume_queue_name()
) )
logging_info("RabbitMQClient register route: [%s]", route) logging_info("RabbitMQClient register route: [%s]", route)
queue: AbstractRobustQueue = await self.channel.declare_queue( if self.amq_configuration.amq_adapter.callbacks_enabled:
name=queue_name, queue: Optional[AbstractRobustQueue] = await self.channel.declare_queue(
durable=True, name=queue_name,
exclusive=False, durable=True,
auto_delete=False, exclusive=False,
arguments=queue_args, auto_delete=False,
) arguments=queue_args,
_exchange_name: str = (
route.exchange if route.exchange else AMQ_RPC_EXCHANGE
)
exchange: AbstractRobustExchange = (
await self.channel.declare_exchange(
name=_exchange_name, type=ExchangeType.topic
) )
else:
queue = None
_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) if self.amq_configuration.amq_adapter.callbacks_enabled:
_c_tag = await queue.consume(callback, no_ack=False) await queue.bind(exchange, route.key)
_c_tag = await queue.consume(callback, no_ack=False)
else:
_c_tag = None
# Register the route with the router and publish its detail to other adapters # Register the route with the router and publish its detail to other adapters
await self.router.register_route( await self.router.register_route(
AMQRoute( AMQRoute(
@@ -155,6 +151,9 @@ class RabbitMQClient:
Set up the RabbitMQ connection Set up the RabbitMQ connection
:return: nothing, it applies changes to instance variables :return: nothing, it applies changes to instance variables
""" """
logging.info(
f"RabbitMQClient.setupAMQ, connecting to RabbitMQ at {self.amq_configuration.dispatch.amq_host}:{self.amq_configuration.dispatch.amq_port}, loop={loop}"
)
self.connection = await aio_pika.connect_robust( self.connection = await aio_pika.connect_robust(
host=self.amq_configuration.dispatch.amq_host, host=self.amq_configuration.dispatch.amq_host,
port=self.amq_configuration.dispatch.amq_port, port=self.amq_configuration.dispatch.amq_port,
@@ -200,16 +199,14 @@ class RabbitMQClient:
message message
), # context path is a routing key ), # context path is a routing key
) )
logging_info( logging_info(f"Msg Publish (no-reply), messageId: {message.id()}, to: {route}")
f"Msg Publish (no-reply), messageId: {message.id}, to: {route}"
)
else: else:
# This is RPC call, there should be response # This is RPC call, there should be response
pika_message: aio_pika.message.Message = aio_pika.message.Message( pika_message: aio_pika.message.Message = aio_pika.message.Message(
body=DataMessageFactory.serialize(message), body=DataMessageFactory.serialize(message),
content_encoding="utf-8", content_encoding="utf-8",
delivery_mode=2, delivery_mode=2,
correlation_id=str(message.id), correlation_id=str(message.id()),
reply_to=self.router.get_reply_to_queue_name(), reply_to=self.router.get_reply_to_queue_name(),
) )
await self.rpc_exchange.publish( await self.rpc_exchange.publish(
@@ -220,12 +217,14 @@ class RabbitMQClient:
) )
logging_info( logging_info(
"Msg Publish (RPC call), messageId: %s, to: %s", "Msg Publish (RPC call), messageId: %s, to: %s",
message.id.as_string(), message.id().as_string(),
route.as_string(), route.as_string(),
) )
else: else:
logging_warning( 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: 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 import re
from typing import Set from typing import Optional, Set
from amqp.adapter.logging_utils import logging_info from amqp.adapter.logging_utils import logging_info
from amqp.model.model import AMQRoute from amqp.model.model import AMQRoute
@@ -41,13 +40,9 @@ class RouteDatabase:
return re.compile(".*") return re.compile(".*")
return re.compile(regex) return re.compile(regex)
def matches( def matches(self, routing_key_from_route: str, routing_key_from_message: str) -> bool:
self, routing_key_from_route: str, routing_key_from_message: str
) -> bool:
route_pattern = self.pattern(routing_key_from_route) route_pattern = self.pattern(routing_key_from_route)
return bool( return bool(route_pattern.match(sanitize_as_rabbitmq_name(routing_key_from_message)))
route_pattern.match(sanitize_as_rabbitmq_name(routing_key_from_message))
)
def wrap_routes(self) -> str: def wrap_routes(self) -> str:
return ",".join(str(route) for route in self.routes) return ",".join(str(route) for route in self.routes)
@@ -60,6 +55,13 @@ class RouteDatabase:
return route return route
return None 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: def add_route(self, route: AMQRoute) -> bool:
if route not in self.routes: if route not in self.routes:
self.routes.add(route) self.routes.add(route)
+17 -20
View File
@@ -1,8 +1,7 @@
import logging from typing import Callable, List, Optional, Set
from typing import Callable, List, Set
from amqp.adapter.amq_route_factory import AMQRouteFactory, NULL_ROUTE from amqp.adapter.amq_route_factory import NULL_ROUTE, AMQRouteFactory
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.model.model import AMQRoute, DataMessage from amqp.model.model import AMQRoute, DataMessage
from amqp.router.route_database import RouteDatabase from amqp.router.route_database import RouteDatabase
from amqp.router.utils import sanitize_as_rabbitmq_name from amqp.router.utils import sanitize_as_rabbitmq_name
@@ -28,7 +27,7 @@ class RouterBase:
self.route_removed_notifier: Callable[[AMQRoute], None] | None = None self.route_removed_notifier: Callable[[AMQRoute], None] | None = None
def get_routing_key(self, message: DataMessage) -> str: 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: def wrap_routes(self) -> str:
return self.route_database.wrap_routes() return self.route_database.wrap_routes()
@@ -38,10 +37,7 @@ class RouterBase:
def unwrap_route_list(self, route_string: str) -> List[AMQRoute]: def unwrap_route_list(self, route_string: str) -> List[AMQRoute]:
logging_debug("RouterMaster.unwrapRoutes, route(s): %s", route_string) logging_debug("RouterMaster.unwrapRoutes, route(s): %s", route_string)
routes = [ routes = [AMQRouteFactory.from_string(route_str) for route_str in route_string.split(",")]
AMQRouteFactory.from_string(route_str)
for route_str in route_string.split(",")
]
return [route for route in routes if route != NULL_ROUTE] return [route for route in routes if route != NULL_ROUTE]
def get_routes(self) -> Set[AMQRoute]: def get_routes(self) -> Set[AMQRoute]:
@@ -51,14 +47,20 @@ class RouterBase:
return self.route_database.find_route(domain, path) return self.route_database.find_route(domain, path)
def find_route_by_message(self, message: DataMessage) -> AMQRoute | None: 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]:
_route: Optional[AMQRoute] = self.route_database.find_route_by_service(service_name)
if _route is None:
for route in self.own_routes:
if route.component_name == service_name:
_route = route
return _route
def add_routes(self, message: str) -> None: def add_routes(self, message: str) -> None:
try: try:
amq_route_list = self.unwrap_route_list(message) amq_route_list = self.unwrap_route_list(message)
logging_info( logging_info(" [DBG] RouterMaster.addRoute [%d]: %s", len(amq_route_list), message)
" [DBG] RouterMaster.addRoute [%d]: %s", len(amq_route_list), message
)
for route in amq_route_list: for route in amq_route_list:
logging_info("RouterMaster.addRoute, adding route: %s", route) logging_info("RouterMaster.addRoute, adding route: %s", route)
if not route.exchange and not route.queue and not route.key: if not route.exchange and not route.queue and not route.key:
@@ -82,14 +84,9 @@ class RouterBase:
def remove_routes(self, message: str) -> None: def remove_routes(self, message: str) -> None:
try: try:
to_remove = [ to_remove = [AMQRouteFactory.from_string(route_str) for route_str in message.split(",")]
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)) removed_count = sum(1 for route in to_remove if self.remove_route(route))
logging_info( logging_info(f"RouterMaster.removeRoutes(): {message}, removed={removed_count}")
f"RouterMaster.removeRoutes(): {message}, removed={removed_count}"
)
except IOError as err: except IOError as err:
logging_error(f"Can't remove route(s): {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 typing import Set
from aio_pika.abc import ( from aio_pika.abc import (
AbstractIncomingMessage,
AbstractRobustChannel, AbstractRobustChannel,
AbstractRobustQueue, AbstractRobustQueue,
AbstractIncomingMessage,
) )
from amqp.adapter.amq_route_factory import AMQRouteFactory from amqp.adapter.amq_route_factory import AMQRouteFactory
from amqp.adapter.logging_utils import ( from amqp.adapter.logging_utils import (
logging_debug,
logging_error, logging_error,
logging_info, logging_info,
logging_debug,
logging_warning, logging_warning,
) )
from amqp.config.amq_configuration import AMQConfiguration 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.model.service_message_type import ServiceMessageType
from amqp.router.router_base import ( from amqp.router.router_base import (
ANY_RECIPIENT,
CM_C_REQ_QUEUE, CM_C_REQ_QUEUE,
CM_REQUEST_EXCHANGE, CM_REQUEST_EXCHANGE,
CM_RESPONSE_EXCHANGE, CM_RESPONSE_EXCHANGE,
ANY_RECIPIENT,
) )
from amqp.router.router_producer import RouterProducer from amqp.router.router_producer import RouterProducer
@@ -77,17 +77,13 @@ class RouterConsumer(RouterProducer):
except Exception as ioe: except Exception as ioe:
logging_error("RouterConsumer FATAL: cannot initialize, %s", ioe) logging_error("RouterConsumer FATAL: cannot initialize, %s", ioe)
else: else:
logging_error( logging_error("RouterConsumer FATAL: cannot initialize, RabbitMQ channel is not open.")
"RouterConsumer FATAL: cannot initialize, RabbitMQ channel is not open."
)
""" """
* ServiceMessageType.ROUTING_DATA_REQUEST handler listening on queue 'CleverMicroRequest' * ServiceMessageType.ROUTING_DATA_REQUEST handler listening on queue 'CleverMicroRequest'
""" """
async def handle_routing_data_request_message( async def handle_routing_data_request_message(self, message: AbstractIncomingMessage):
self, message: AbstractIncomingMessage
):
# some debug statements, TODO - remove for prod release # some debug statements, TODO - remove for prod release
delivery_tag = message.delivery_tag delivery_tag = message.delivery_tag
debug = "Delivery.Properties = " + message.properties.__str__() debug = "Delivery.Properties = " + message.properties.__str__()
@@ -101,9 +97,7 @@ class RouterConsumer(RouterProducer):
await message.ack(False) await message.ack(False)
cm_message: ServiceMessage = self.service_message_factory.from_bytes( cm_message: ServiceMessage = self.service_message_factory.from_bytes(message.body)
message.body
)
route_mapping: str = self.wrap_own_routes() route_mapping: str = self.wrap_own_routes()
logging_debug( logging_debug(
"******[%s] (ROUTING_DATA_REQ) msg=%s, routeMap=[%s]", "******[%s] (ROUTING_DATA_REQ) msg=%s, routeMap=[%s]",
@@ -124,9 +118,7 @@ class RouterConsumer(RouterProducer):
route_mapping, route_mapping,
cm_message.reply_to, cm_message.reply_to,
) )
await self.publish_service_message( await self.publish_service_message(service_message, self.cm_response_exchange)
service_message, self.cm_response_exchange
)
except Exception as err: except Exception as err:
logging_error( logging_error(
"******[%s] (ROUTING_DATA_REQ) Response to queue: %s failed: %s", "******[%s] (ROUTING_DATA_REQ) Response to queue: %s failed: %s",
@@ -158,9 +150,7 @@ class RouterConsumer(RouterProducer):
serviceMessage: ServiceMessage = self.service_message_factory.of( serviceMessage: ServiceMessage = self.service_message_factory.of(
ServiceMessageType.ROUTING_DATA, route.as_string(), ANY_RECIPIENT ServiceMessageType.ROUTING_DATA, route.as_string(), ANY_RECIPIENT
) )
await self.publish_service_message( await self.publish_service_message(serviceMessage, self.cm_response_exchange)
serviceMessage, self.cm_response_exchange
)
self.add_own_route(route) self.add_own_route(route)
logging_info("RouterConsumer registered Route: %s", route) logging_info("RouterConsumer registered Route: %s", route)
else: else:
@@ -182,16 +172,14 @@ class RouterConsumer(RouterProducer):
service_message: ServiceMessage = self.service_message_factory.of( service_message: ServiceMessage = self.service_message_factory.of(
ServiceMessageType.ROUTING_DATA_REMOVE, route.as_string(), ANY_RECIPIENT ServiceMessageType.ROUTING_DATA_REMOVE, route.as_string(), ANY_RECIPIENT
) )
if not await self.publish_service_message( if not await self.publish_service_message(service_message, self.cm_response_exchange):
service_message, self.cm_response_exchange
):
logging_warning( logging_warning(
"RouterConsumer couldn't remove current route: %s, channel not open.", "RouterConsumer couldn't remove current_availability route: %s, channel not open.",
route, route,
) )
except Exception as ioe: 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) result: bool = self.remove_route(route)
return result return result
@@ -201,13 +189,9 @@ class RouterConsumer(RouterProducer):
* Cleanup - de-register routes with master. * Cleanup - de-register routes with master.
""" """
routes: Set[AMQRoute] = self.get_routes() routes: Set[AMQRoute] = self.get_routes()
logging_info( logging_info("RouterConsumer.cleanup() [%d]: %s", len(routes), self.wrap_routes())
"RouterConsumer.cleanup() [%d]: %s", len(routes), self.wrap_routes()
)
not_removed = sum( not_removed = sum(
1 1 for r in (self.remove_route_on_cleanup(route) for route in routes) if not r
for r in (self.remove_route_on_cleanup(route) for route in routes)
if not r
) )
if not_removed > 0: if not_removed > 0:
logging_warning( logging_warning(
+17 -30
View File
@@ -7,22 +7,20 @@
* forwards it to the application. * forwards it to the application.
""" """
import logging
from aio_pika import Message from aio_pika import Message
from aio_pika.abc import ( from aio_pika.abc import (
AbstractIncomingMessage,
AbstractRobustChannel, AbstractRobustChannel,
AbstractRobustQueue,
AbstractRobustConnection, AbstractRobustConnection,
AbstractRobustExchange, AbstractRobustExchange,
AbstractIncomingMessage, AbstractRobustQueue,
) )
from pika.exchange_type import ExchangeType from pika.exchange_type import ExchangeType
from amqp.adapter.logging_utils import ( from amqp.adapter.logging_utils import (
logging_debug,
logging_error, logging_error,
logging_info, logging_info,
logging_debug,
logging_warning, logging_warning,
) )
from amqp.adapter.service_message_factory import ServiceMessageFactory 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.model import ServiceMessage
from amqp.model.service_message_type import ServiceMessageType from amqp.model.service_message_type import ServiceMessageType
from amqp.router.router_base import ( 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_REQUEST_EXCHANGE,
CM_RESPONSE_EXCHANGE, CM_RESPONSE_EXCHANGE,
CM_C_AMQ_QUEUE, RouterBase,
CM_P_RESP_QUEUE,
CM_P_REPLY_TO,
ANY_RECIPIENT,
) )
@@ -49,16 +47,14 @@ class RouterProducer(RouterBase):
""" """
super().__init__() super().__init__()
self.connection: AbstractRobustConnection | None = None self.connection: AbstractRobustConnection | None = None
self.service_message_factory = ServiceMessageFactory( self.service_message_factory = ServiceMessageFactory(configuration.amq_adapter.service_name)
configuration.amq_adapter.service_name
)
self.amq_configuration: AMQConfiguration = configuration self.amq_configuration: AMQConfiguration = configuration
self.channel: AbstractRobustChannel | None = None self.channel: AbstractRobustChannel | None = None
self.cm_request_exchange: AbstractRobustExchange | None = None self.cm_request_exchange: AbstractRobustExchange | None = None
self.cm_response_exchange: AbstractRobustExchange | None = None self.cm_response_exchange: AbstractRobustExchange | None = None
logging_info( 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.service_name,
self.amq_configuration.amq_adapter.route_mapping, self.amq_configuration.amq_adapter.route_mapping,
) )
@@ -85,7 +81,7 @@ class RouterProducer(RouterBase):
name=CM_REQUEST_EXCHANGE, type=ExchangeType.fanout 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) # a RoutingData request)
self.cm_response_exchange = await self.channel.declare_exchange( self.cm_response_exchange = await self.channel.declare_exchange(
name=CM_RESPONSE_EXCHANGE, type=ExchangeType.fanout name=CM_RESPONSE_EXCHANGE, type=ExchangeType.fanout
@@ -100,9 +96,7 @@ class RouterProducer(RouterBase):
arguments={}, arguments={},
) )
# 2b. bind queue to CM_RESPONSE_EXCHANGE w/o any conditions (empty routing key) # 2b. bind queue to CM_RESPONSE_EXCHANGE w/o any conditions (empty routing key)
await _response_queue.bind( await _response_queue.bind(exchange=CM_RESPONSE_EXCHANGE, routing_key="#")
exchange=CM_RESPONSE_EXCHANGE, routing_key="#"
)
# 2c. listen for incoming RoutingData messages # 2c. listen for incoming RoutingData messages
_c_tag = await _response_queue.consume( _c_tag = await _response_queue.consume(
callback=self.handle_routing_data_message, no_ack=False callback=self.handle_routing_data_message, no_ack=False
@@ -116,9 +110,7 @@ class RouterProducer(RouterBase):
logging_error("RouterProducer FATAL: cannot initialize, %s", ioe) logging_error("RouterProducer FATAL: cannot initialize, %s", ioe)
else: else:
logging_error( logging_error("RouterProducer FATAL: cannot initialize, RabbitMQ channel is not open.")
"RouterProducer FATAL: cannot initialize, RabbitMQ channel is not open."
)
async def handle_routing_data_message(self, message: AbstractIncomingMessage): 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. * routing data, invokes addRoute or removeRoute depending on the type of the message.
""" """
logging_info( logging_info(
f"[%s] Received ROUTING DATA: [%s], msgId=%s", "[%s] Received ROUTING DATA: [%s], msgId=%s",
message.delivery_tag, message.delivery_tag,
message.body, message.body,
message.message_id, message.message_id,
) )
await message.ack(multiple=False) await message.ack(multiple=False)
cm_message: ServiceMessage = self.service_message_factory.from_bytes( cm_message: ServiceMessage = self.service_message_factory.from_bytes(message.body)
message.body
)
# ignore the message if it comes from ourselves (even if from different instance) # ignore the message if it comes from ourselves (even if from different instance)
logging_debug( logging_debug(
"******[%s] ROUTING_DATA, msg=%s, routeMap=[%s] sn=%s, replyTo=%s", "******[%s] ROUTING_DATA, msg=%s, routeMap=[%s] sn=%s, replyTo=%s",
@@ -146,8 +136,7 @@ class RouterProducer(RouterBase):
) )
if ( if (
cm_message.is_valid() cm_message.is_valid()
and not self.amq_configuration.amq_adapter.service_name and not self.amq_configuration.amq_adapter.service_name == cm_message.reply_to
== cm_message.reply_to
): ):
if cm_message.message_type == ServiceMessageType.ROUTING_DATA: if cm_message.message_type == ServiceMessageType.ROUTING_DATA:
self.add_routes(cm_message.message) self.add_routes(cm_message.message)
@@ -166,9 +155,7 @@ class RouterProducer(RouterBase):
serviceMessage: ServiceMessage = self.service_message_factory.of( serviceMessage: ServiceMessage = self.service_message_factory.of(
ServiceMessageType.ROUTING_DATA_REQUEST, "", ANY_RECIPIENT ServiceMessageType.ROUTING_DATA_REQUEST, "", ANY_RECIPIENT
) )
if not await self.publish_service_message( if not await self.publish_service_message(serviceMessage, self.cm_request_exchange):
serviceMessage, self.cm_request_exchange
):
logging_warning( logging_warning(
"RouterProducer could not request route data: %s channel is not open.", "RouterProducer could not request route data: %s channel is not open.",
CM_REQUEST_EXCHANGE, CM_REQUEST_EXCHANGE,
+293 -57
View File
@@ -1,32 +1,36 @@
import asyncio import asyncio
import importlib
import inspect
import os
import threading import threading
import time import time
import os from asyncio import AbstractEventLoop, Future
from asyncio import Future, AbstractEventLoop from typing import Any, Dict
from typing import Dict
from aio_pika import Message from aio_pika import Message
from aio_pika.abc import AbstractIncomingMessage, AbstractRobustExchange from aio_pika.abc import AbstractIncomingMessage, AbstractRobustExchange
from opentelemetry import trace from opentelemetry import trace
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from opentelemetry.trace import Tracer, TraceState 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.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_message_factory import DataMessageFactory
from amqp.adapter.data_response_factory import DataResponseFactory from amqp.adapter.data_response_factory import DataResponseFactory
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter, AMQMessage from amqp.adapter.file_handler import FileHandler, StreamingFileHandler
from amqp.adapter.file_handler import StreamingFileHandler, FileHandler from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
from amqp.adapter.logging_utils import logging_error, logging_debug, logging_info from amqp.adapter.serializer import map_as_string, parse_map_string
from amqp.adapter.service_message_factory import ServiceMessageFactory from amqp.adapter.service_message_factory import ServiceMessageFactory
from amqp.config.amq_configuration import AMQConfiguration from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import ( from amqp.model.model import (
DataResponse,
DataMessage, DataMessage,
DataMessageMagicByte,
DataResponse,
MultipartDataMessage, MultipartDataMessage,
) )
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient
from amqp.adapter.serializer import map_as_string from amqp.router.utils import sanitize_as_rabbitmq_name
def collect_trace_info(): def collect_trace_info():
@@ -45,6 +49,75 @@ def set_text(carrier, key, value):
carrier[key] = value carrier[key] = value
async def invoke_command(
package_name: str, class_name: str, function_name: str, msg_id: str = "", kwargs=None
) -> Any:
if kwargs is None:
kwargs = {}
try:
# 1. Dynamically import the package
logging_info(f"Attempting to import package: {package_name}")
package = importlib.import_module(package_name)
logging_info(f"Package '{package_name}' imported successfully.")
target_callable = None
if class_name == "__global__":
# 2. Handle global function invocation
logging_info(f"Attempting to find global function: {function_name} in {package_name}")
target_callable = getattr(package, function_name)
logging_info(f"Global function '{function_name}' found.")
else:
# 3. Handle class method invocation
logging_info(f"Attempting to find class: {class_name} in {package_name}")
target_class = getattr(package, class_name)
logging_info(f"Class '{class_name}' found. Instantiating...")
# Instantiate the class. Assuming a default constructor or one that doesn't
# require arguments from kwargs for initialization.
# If the class __init__ needs specific arguments, this part might need adjustment
# based on how those arguments are provided.
instance = target_class()
logging_info(f"Instance of '{class_name}' created.")
logging_info(f"Attempting to find method: {function_name} in class {class_name}")
target_callable = getattr(instance, function_name)
logging_info(f"Method '{function_name}' found.")
# 4. Invoke the function/method
logging_info(f"Invoking '{function_name}' with arguments: {kwargs}")
if inspect.iscoroutinefunction(target_callable):
result = await target_callable(**kwargs)
logging_info(f"Asynchronous function '{function_name}' invoked. Result: {result}")
else:
result = target_callable(**kwargs)
logging_info(f"Synchronous function '{function_name}' invoked. Result: {result}")
response: DataResponse = DataResponseFactory.of(
id=msg_id,
response_code=200,
content_type="text/plain",
body=result.encode("utf-8"),
trace_info=collect_trace_info(),
error=None,
error_cause=None,
)
response._magic = DataMessageMagicByte.RPC_REQUEST.value
return response
except ModuleNotFoundError as e:
logging_error(f"Error: Package '{package_name}' not found. {e}")
raise
except AttributeError as e:
logging_error(
f"Error: Class '{class_name}' or function/method '{function_name}' not found in '{package_name}'. {e}"
)
raise
except Exception as e:
logging_error(f"An unexpected error occurred during invocation: {e}")
raise
class DataMessageHandler: class DataMessageHandler:
def __init__( def __init__(
self, self,
@@ -75,21 +148,15 @@ class DataMessageHandler:
self.file_handler: FileHandler = file_handler self.file_handler: FileHandler = file_handler
self.backpressure_handler: BackpressureHandler = backpressure_handler self.backpressure_handler: BackpressureHandler = backpressure_handler
async def inbound_data_message_callback_as_thread( async def inbound_data_message_callback_as_thread(self, message: AbstractIncomingMessage):
self, message: AbstractIncomingMessage
):
""" """
The callback method that is called when a DataMessage is received from RabbitMQ, starts new Thread. The callback method that is called when a DataMessage is received from RabbitMQ, starts new Thread.
""" """
threading.Thread( threading.Thread(
target=lambda: asyncio.run( target=lambda: asyncio.run(self.inbound_data_message_callback_no_thread(message))
self.inbound_data_message_callback_no_thread(message)
)
).start() ).start()
async def inbound_data_message_callback_no_thread( async def inbound_data_message_callback_no_thread(self, message: AbstractIncomingMessage):
self, message: AbstractIncomingMessage
):
""" """
Receives the DataMessage as bytes from RabbitMQ queue and deserializes it to the DataMessage object Receives the DataMessage as bytes from RabbitMQ queue and deserializes it to the DataMessage object
Ensures Jaeger trace-id propagation. Ensures Jaeger trace-id propagation.
@@ -101,19 +168,19 @@ class DataMessageHandler:
amq_message = await self.reconstitute_data_message(message) amq_message = await self.reconstitute_data_message(message)
if amq_message is not None: if amq_message is not None:
self.ensure_trace_info(amq_message) 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) await self.run_http_request_magic(message, amq_message, self.loop)
elif amq_message.magic() == DataMessageMagicByte.RPC_REQUEST.value:
await self.run_rpc_request(message, amq_message, self.loop)
else: else:
logging_error( logging_error(
"Unknown MAGIC value: [%s], don't know how to interpret the message", "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) asyncio.run_coroutine_threadsafe(message.nack(requeue=False), self.loop)
return return
else: else:
logging_error( logging_error("######[%s] No valid AMQ Message present.", message.delivery_tag)
"######[%s] No valid AMQ Message present.", message.delivery_tag
)
self.record_duration(start, message.delivery_tag) self.record_duration(start, message.delivery_tag)
asyncio.run_coroutine_threadsafe(message.ack(multiple=False), self.loop) asyncio.run_coroutine_threadsafe(message.ack(multiple=False), self.loop)
@@ -131,16 +198,16 @@ class DataMessageHandler:
:param loop: Event loop :param loop: Event loop
:return: DataResponse :return: DataResponse
""" """
# Increment the counter for parallel executions # Increment the counter for parallel executions and check resource availability
self.backpressure_handler.increase_parallel_executions() self.backpressure_handler.increase_current_load()
await self.backpressure_handler.check_overload_condition()
logging_info( 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: if message.reply_to:
self.reply_to = message.reply_to self.reply_to = message.reply_to
try: try:
await self.backpressure_handler.check_overload_condition()
_a_message = AMQMessage(amq_message, self.rabbit_mq_client.connection) _a_message = AMQMessage(amq_message, self.rabbit_mq_client.connection)
_response: DataResponse = await self.service_adapter.on_message(_a_message) _response: DataResponse = await self.service_adapter.on_message(_a_message)
logging_debug(f"Result in thread: {_response}") logging_debug(f"Result in thread: {_response}")
@@ -163,11 +230,96 @@ class DataMessageHandler:
except Exception as e: except Exception as e:
logging_error(f"Request handler: {e}") logging_error(f"Request handler: {e}")
self.backpressure_handler.decrease_parallel_executions() self.backpressure_handler.decrease_current_load()
async def run_rpc_request(
self,
message: AbstractIncomingMessage,
amq_message: DataMessage,
loop: AbstractEventLoop,
) -> None:
"""
Synchronous version of on_message method. This is a wrapper around the async on_message method.
It also counts parallel executions and handles backpressure.
:param message: RabbitMQ raw message
:param amq_message: DataMessage
:param loop: Event loop
:return: DataResponse
"""
# 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_availability}] of [{self.backpressure_handler.max_availability}]"
)
if message.reply_to:
self.reply_to = message.reply_to
try:
_a_message = AMQMessage(amq_message, self.rabbit_mq_client.connection)
_response: DataResponse = await self.on_rpc(_a_message)
logging_debug(f"Result in thread: {_response}")
try:
await self.send_reply(message.reply_to, _response)
except Exception as e:
logging_error(f"Processing message: {e}")
asyncio.run_coroutine_threadsafe(message.nack(requeue=False), loop=loop)
if isinstance(amq_message, MultipartDataMessage):
# remove the downloaded files from the output_dir
for file_id, filename in amq_message.extra_data.items():
try:
logging_info(
f"[{message.delivery_tag}] Deleting file: {file_id} in {filename}"
)
os.remove(filename)
except OSError as e:
logging_error(f"Deleting file {filename}: {e}")
except Exception as e:
logging_error(f"Request handler: {e}")
self.backpressure_handler.decrease_current_load()
async def on_rpc(self, a_message: AMQMessage) -> DataResponse:
"""
Dynamically invokes a function or a class method from a specified package.
Args:
a_message (AMQMessage): The message containing the package, class, and function details.
Args encoded in the AMQMessage:
package_name (str): The name of the Python package to import.
class_name (str): The name of the class within the package.
Use '__global__' if the function_name refers to a global function
in the package.
function_name (str): The name of the function or method to invoke.
kwargs (Dict[str, Any]): A dictionary of keyword arguments to pass to the function/method (in the body).
Returns:
Any: The return value of the invoked function or method.
Raises:
ModuleNotFoundError: If the specified package cannot be imported.
AttributeError: If the specified class or function/method does not exist.
Exception: For any other errors during invocation.
"""
package_name: str = a_message.method()
class_name: str = a_message.domain()
function_name: str = a_message.path()
kwargs: Dict[str, Any] = parse_map_string(a_message._base64body.decode("utf-8"))
return await invoke_command(
package_name, class_name, function_name, a_message.id().as_string(), kwargs
)
async def reconstitute_data_message( async def reconstitute_data_message(
self, message: AbstractIncomingMessage self, message: AbstractIncomingMessage
) -> DataMessage | None: ) -> 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 amq_message: DataMessage | None = None
delivery_tag = message.delivery_tag delivery_tag = message.delivery_tag
addressing: int = message.headers.get("clevermicro.addressing", 0) addressing: int = message.headers.get("clevermicro.addressing", 0)
@@ -177,9 +329,7 @@ class DataMessageHandler:
amq_message = await DataMessageFactory.from_stream( amq_message = await DataMessageFactory.from_stream(
message.body, rabbitmq_url=self.rabbit_mq_url, loop=self.loop message.body, rabbitmq_url=self.rabbit_mq_url, loop=self.loop
) )
logging_info( logging_info(f"######[{delivery_tag}] AMQ Message from stream: {amq_message}")
f"######[{delivery_tag}] AMQ Message from stream: {amq_message}"
)
if amq_message is not None: if amq_message is not None:
extra_data = await self.file_handler.on_message( extra_data = await self.file_handler.on_message(
message=message, message=message,
@@ -188,26 +338,24 @@ class DataMessageHandler:
loop=self.loop, loop=self.loop,
output_dir=self.output_dir, output_dir=self.output_dir,
) )
logging_info( logging_info(f"######[{delivery_tag}] Multipart download result: {extra_data}")
f"######[{delivery_tag}] Multipart download result: {extra_data}"
)
amq_message = MultipartDataMessage(amq_message, extra_data) amq_message = MultipartDataMessage(amq_message, extra_data)
if amq_message: if amq_message:
logging_info( logging_info(
"######[%s] AMQ Message received, ctag=%s, msgId=%s, %s:%s, traceInfo=%s", "######[%s] AMQ Message received, ctag=%s, msgId=%s, %s:%s, traceInfo=%s",
message.delivery_tag, message.delivery_tag,
message.consumer_tag, message.consumer_tag,
amq_message.id, amq_message.id(),
amq_message.method, amq_message.method(),
amq_message.path, amq_message.path(),
map_as_string(amq_message.trace_info), map_as_string(amq_message.trace_info()),
) )
return amq_message return amq_message
def ensure_trace_info(self, amq_message: DataMessage): def ensure_trace_info(self, amq_message: DataMessage):
propagator = TraceContextTextMapPropagator() propagator = TraceContextTextMapPropagator()
context = propagator.extract( 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) self.tracer.start_span(name="AMQ-Adapter-Consumer", context=context)
# logging_info("CTX=%s", context) # logging_info("CTX=%s", context)
@@ -216,15 +364,15 @@ class DataMessageHandler:
# logging.info(" [*] CTX2=%s", context_with_span) # logging.info(" [*] CTX2=%s", context_with_span)
# propagator.inject(context_with_span, amq_message.trace_info, self.set_text) # propagator.inject(context_with_span, amq_message.trace_info, self.set_text)
# logging.info(" [***] NEW TI=%s CTX=%s SCOPE=%s", # 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): def record_duration(self, start, delivery_tag):
global last_data_message_time global last_backpressure_event_time
last_data_message_time = time.time() last_backpressure_event_time = time.time()
logging_info( logging_info(
"######[%s] Done, single ACK: %sms.", "######[%s] Done, single ACK: %sms.",
delivery_tag, delivery_tag,
last_data_message_time - start, last_backpressure_event_time - start,
) )
async def send_reply(self, reply_to: str, response: DataResponse): async def send_reply(self, reply_to: str, response: DataResponse):
@@ -238,17 +386,15 @@ class DataMessageHandler:
if reply_to: if reply_to:
pika_message: Message = Message( pika_message: Message = Message(
body=DataResponseFactory.serialize(response), body=DataResponseFactory.serialize(response),
correlation_id=response.id, correlation_id=response.id(),
content_type=( content_type=(
response.content_type[0] response.content_type()[0]
if isinstance(response.content_type, list) if isinstance(response.content_type(), list)
else response.content_type else response.content_type()
), ),
) )
_res = asyncio.run_coroutine_threadsafe( _res = asyncio.run_coroutine_threadsafe(
self.reply_to_exchange.publish( self.reply_to_exchange.publish(message=pika_message, routing_key=reply_to),
message=pika_message, routing_key=reply_to
),
self.loop, self.loop,
) )
logging_debug( logging_debug(
@@ -260,21 +406,18 @@ class DataMessageHandler:
async def reply_received_callback(self, message: AbstractIncomingMessage): async def reply_received_callback(self, message: AbstractIncomingMessage):
""" """
The reply for DataMessage that we sent out earlier - must match the received response with The handler for reply for DataMessage that we sent out earlier -
the original request and respond to the caller. the code matches the received response with the original request and passes the response to the caller.
:param message: :param message:
:return: :return:
""" """
reply_id = message.correlation_id reply_id = message.correlation_id
content_type = message.content_type
delivery_tag = message.delivery_tag delivery_tag = message.delivery_tag
debug = f"######[{delivery_tag}] (RPC REPLY) env: {message.headers}, props: {message.properties}, BODY: {message.body.decode()}" debug = f"######[{delivery_tag}] (RPC REPLY) env: {message.headers}, props: {message.properties}, BODY: {message.body.decode()}"
logging_debug(debug) logging_debug(debug)
# TODO - figure out implementation of this part, as this is sending reply back either to a service, # TODO - figure out implementation of this part, as this is sending reply back either to a service,
# or to the user # or to the user
logging_info( logging_info(f"OUTSTANDING RECV: {self.outstanding}, size={len(self.outstanding)}")
f"OUTSTANDING RECV: {self.outstanding}, size={len(self.outstanding)}"
)
future: Future = self.outstanding.pop(reply_id, None) future: Future = self.outstanding.pop(reply_id, None)
# if reply is None: # if reply is None:
@@ -293,3 +436,96 @@ class DataMessageHandler:
future.set_result(response.body()) future.set_result(response.body())
await message.ack(multiple=False) 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=sanitize_as_rabbitmq_name(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
"""
rmq_body = DataMessageFactory.serialize(message)
cid = message.id().as_string()
ctype = (
message.content_type()[0]
if isinstance(message.content_type(), list)
else message.content_type()
)
print(
f"Sending RPC command to route: {route}, message ID: {cid} with body: {rmq_body} correlation ID: {cid}, content type: {ctype}"
)
# Create a Pika message with the correlation ID set to the message ID
pika_message = Message(
body=rmq_body,
correlation_id=cid,
content_type=ctype,
)
print(f"SENDING RPC COMMAND: {pika_message.body} with ID: {message.id().as_string()}")
# 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=sanitize_as_rabbitmq_name(route.component_name)
)
logging_info(
"###### RPC Request Published to exchg:(%s) routing_key=%s, msgId=%s",
rpc_exchange_name,
route.component_name,
message.id().as_string(),
)
return True
+176 -22
View File
@@ -1,25 +1,25 @@
import asyncio
import logging import logging
import logging.config import logging.config
import os import os
import asyncio
from asyncio import Future from asyncio import Future
from threading import Thread from threading import Thread
from typing import Optional from typing import AsyncIterator, Optional
from aio_pika.abc import AbstractRobustExchange 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.amq_route_factory import AMQRouteFactory
from amqp.adapter.logging_utils import logging_info, logging_warning 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_error, logging_info, logging_warning
from amqp.adapter.service_message_factory import ServiceMessageFactory from amqp.adapter.service_message_factory import ServiceMessageFactory
from amqp.config.amq_configuration import AMQConfiguration 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.rabbitmq.rabbit_mq_client import RabbitMQClient
from amqp.router.router_consumer import RouterConsumer
from amqp.service.amq_message_handler import DataMessageHandler 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.basicConfig(level=logging.DEBUG)
logging.info(f"CWD = {os.getcwd()}") logging.info(f"CWD = {os.getcwd()}")
@@ -38,24 +38,44 @@ class AMQService:
def __init__(self, amq_configuration: AMQConfiguration, service_adapter=None): def __init__(self, amq_configuration: AMQConfiguration, service_adapter=None):
logging.info("***********************************************") logging.info("***********************************************")
logging.info("********** AMQ Service *******************") logging.info("************** AMQ Service ***************")
logging.info("***********************************************") logging.info("***********************************************")
self.loop = asyncio.new_event_loop() self.loop = asyncio.new_event_loop()
self.amq_configuration = amq_configuration self.amq_configuration = amq_configuration
self.service_adapter = service_adapter 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.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.amq_data_message_handler = None
self.service_message_factory = ServiceMessageFactory( self.service_message_factory = ServiceMessageFactory(
self.amq_configuration.amq_adapter.service_name self.amq_configuration.amq_adapter.service_name
) )
self.backpressure_thread: Optional[Thread] = None self.backpressure_thread: Optional[Thread] = None
self.backpressure_handler: Optional[BackpressureHandler] = None
self.maximum_availability = -1
self.future: Future = asyncio.Future()
# =======================================================================
# =======================================================================
# ============== P u b l i c A P I M e t h o d s ===================
# =======================================================================
# =======================================================================
def run(self): def run(self):
"""
Start the AMQ service, initializing RabbitMQ connection and message handlers.
"""
# self.init(loop) will initialize RabbitMQ connection # self.init(loop) will initialize RabbitMQ connection
asyncio.set_event_loop(self.loop) asyncio.set_event_loop(self.loop)
self.loop.run_until_complete(self.init(self.loop, self.service_adapter)) self.loop.run_until_complete(self.init(self.loop, self.service_adapter))
@@ -63,13 +83,146 @@ class AMQService:
logging.info("***********************************************") logging.info("***********************************************")
logging.info("****** AMQ Service Init Complete *********") logging.info("****** AMQ Service Init Complete *********")
logging.info("***********************************************") logging.info("***********************************************")
self.future.set_result(True)
self.loop.run_forever() self.loop.run_forever()
async def init(self, loop, service_adapter): def backpressure(self, current_availability: int, maximum: int = -1):
_reply_to_exchange: AbstractRobustExchange = await self.rabbit_mq_client.init( """
loop 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,
)
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 = asyncio.run_coroutine_threadsafe(
self.amq_data_message_handler.send_rpc(route, message),
loop=self.backpressure_handler.loop,
)
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.")
) )
backpressure_handler = BackpressureHandler( return result
def command(self, service_name: str, fq_method_name: str, **kwargs) -> bool:
"""
Send a unidirectional Command 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:
m: DataMessage = self.data_message_factory.create_rpc_message(
fq_method_name=fq_method_name,
swarm_id=self.amq_configuration.amq_adapter.swarm_task_id,
args=kwargs,
)
asyncio.run_coroutine_threadsafe(
self.amq_data_message_handler.send_command(route, m),
loop=self.backpressure_handler.loop,
)
return True
return False
async def data_message_generator(self, route: AMQRoute) -> AsyncIterator[DataMessage]:
"""
An async iterator that blocks until messages arrive, processes them,
and continues in a loop.
:param queue_name: The name of the queue to consume messages from
:yield: The processed DataResponse for each message
"""
# Get a channel from the connection
channel = self.rabbit_mq_client.channel
queue_args = (
self.rabbit_mq_client.DLQ_ARGS if self.amq_configuration.dispatch.use_dlq else None
)
queue = await channel.declare_queue(
route.queue,
durable=True,
exclusive=False,
auto_delete=False,
arguments=queue_args,
)
await queue.bind(route.exchange, route.key)
# Start consuming without a callback
async with queue.iterator() as queue_iter:
async for message in queue_iter:
try:
# Process the message using the same logic as in inbound_data_message_callback_no_thread
data_message: DataMessage = (
await self.amq_data_message_handler.reconstitute_data_message(message)
)
success = yield data_message
if not success:
logging_warning(f"Message processing failed for {data_message.id()}.")
# Nack the message without requeuing in case of failure
await message.nack(requeue=False)
else:
await message.ack()
except Exception as e:
logging_error(f"Error processing message: {e}")
# Nack the message without requeuing in case of error
await message.nack(requeue=False)
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, channel=self.rabbit_mq_client.channel,
loop=loop, loop=loop,
config=self.amq_configuration, config=self.amq_configuration,
@@ -83,11 +236,12 @@ class AMQService:
self.rabbit_mq_client, self.rabbit_mq_client,
self.amq_configuration, self.amq_configuration,
loop=loop, loop=loop,
backpressure_handler=backpressure_handler, backpressure_handler=self.backpressure_handler,
) )
await self.register_routes() await self.register_routes()
await self.rabbit_mq_client.request_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): def cleanup(self):
logging_info("AMQService.cleanup() rmqc=%s", self.rabbit_mq_client) logging_info("AMQService.cleanup() rmqc=%s", self.rabbit_mq_client)
@@ -96,7 +250,7 @@ class AMQService:
async def reinitialize_if_disconnected(self): async def reinitialize_if_disconnected(self):
if not self.rabbit_mq_client.channel or self.rabbit_mq_client.channel.is_closed: 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.rabbit_mq_client.init(loop=self.loop)
await self.register_routes() await self.register_routes()
await self.rabbit_mq_client.request_routes() await self.rabbit_mq_client.request_routes()
@@ -121,8 +275,8 @@ class AMQService:
) )
return None return None
def send_message(self, message: DataMessage, future: Future): def set_outstanding_future(self, message: DataMessage, future: Future):
self.amq_data_message_handler.outstanding[str(message.id)] = future self.amq_data_message_handler.outstanding[str(message.id())] = future
def remove_outstanding_future(self, message_id: str): def remove_outstanding_future(self, message_id: str):
self.amq_data_message_handler.outstanding.pop(message_id, None) self.amq_data_message_handler.outstanding.pop(message_id, None)
+186 -92
View File
@@ -1,102 +1,196 @@
from opentelemetry import metrics, trace """
from opentelemetry.sdk.metrics import MeterProvider OpenTelemetry configuration for tracing and metrics.
from opentelemetry.sdk.trace import TracerProvider, Tracer
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from prometheus_client import start_http_server This module provides functionality to initialize and configure OpenTelemetry
from opentelemetry.exporter.prometheus import PrometheusMetricReader for both tracing (Jaeger via OTLP) and metrics (Prometheus) in the AMQ adapter.
from opentelemetry.sdk.resources import SERVICE_NAME, Resource """
import logging
import os 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: @dataclass
# Resource can be required for some backends, e.g. Jaeger or Prometheus class TracingConfig:
# If resource wouldn't be set - traces wouldn't appear in Jaeger. """Configuration for OpenTelemetry tracing."""
# The 'name' value is the name shown in Jaeger.
resource = Resource( service_name: str
attributes={ service_version: str
SERVICE_NAME: ( otlp_endpoint: str = "http://localhost:4317" # Default Jaeger OTLP endpoint
name if name else os.environ.get("APPLICATION_NAME", "Python App") insecure: bool = True
) debug: bool = False
} additional_attributes: Optional[Dict[str, Any]] = None
)
#
# 1. Tracing (Jaeger) configuration @dataclass
# ---------------------------------- class MetricsConfig:
# """Configuration for OpenTelemetry metrics."""
# Configure the OTLP exporter (or HTTP exporter the same way)
# the recommended configuration is via environment variables. The mandatory ones: prometheus_port: int = 9464
# OTEL_EXPORTER_OTLP_ENDPOINT: "http://jaeger:4317" prometheus_host: str = "localhost"
# OTEL_EXPORTER_OTLP_INSECURE: "true" export_interval_millis: int = 30000
# for more see here: https://opentelemetry-python.readthedocs.io/en/latest/exporter/otlp/otlp.html additional_attributes: Optional[Dict[str, Any]] = None
# alternatively these values can be provided as input values to OTLPSpanExporter() constructor.
#
# otlp_exporter = OTLPSpanExporter() def create_resource(config: TracingConfig) -> Resource:
# or provide endpoint explicitly, as """Create an OpenTelemetry resource with service information."""
otlp_exporter = OTLPSpanExporter(endpoint="http://jaeger:4317") attributes = {
# SERVICE_NAME: config.service_name,
# Initialize the tracer provider. This is abstract API call. SERVICE_VERSION: config.service_version,
# 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) provider = TracerProvider(resource=resource)
# Add the exporters to the batch span processor
span_processor = BatchSpanProcessor(otlp_exporter)
provider.add_span_processor(span_processor)
# Sets the global default tracer provider
trace.set_tracer_provider(provider)
#
# logging.info(' * Instrumenting application..... ') # Add OTLP exporter for Jaeger
# if Flask, then use: otlp_exporter = setup_otlp_exporter(config)
# FlaskInstrumentor().instrument_app(flask_app) provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
# if FastAPI, then use
# FastAPIInstrumentor.instrument_app(app)
# # Add console exporter in debug mode
# Python Flask is now automatically instrumented to export details of each REST endpoint call. if config.debug:
# For Django based app, use 'opentelemetry-instrumentation-django' import instead. see here: provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))
# 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.
# 2. Metrics (Prometheus) configuration return provider
# -------------------------------------
#
# Start Prometheus client
start_http_server(port=9464, addr="localhost")
# Initialize PrometheusMetricReader which pulls metrics from the SDK
# on-demand to respond to scrape requests (unlike Jaeger, which worls with push model - application
# pushes traces when those become available to Jaeger's collector. Prometheus works with pull model
# where it periodically polls pre-configured and/or discovered endpoints for metrics.
reader = PrometheusMetricReader()
provider = MeterProvider(resource=resource, metric_readers=[reader])
metrics.set_meter_provider(provider)
return trace.get_tracer(name)
def setup_metrics_provider(config: MetricsConfig, resource: Resource) -> MeterProvider:
"""Set up the MeterProvider with Prometheus export."""
try:
# Start Prometheus HTTP server
start_http_server(port=config.prometheus_port, addr=config.prometheus_host)
logger.info(
f"Started Prometheus metrics server on {config.prometheus_host}:{config.prometheus_port}"
)
# Create Prometheus reader
prometheus_reader = PrometheusMetricReader()
# Create periodic reader for custom metrics
periodic_reader = PeriodicExportingMetricReader(
prometheus_reader, export_interval_millis=config.export_interval_millis
)
# Create and configure meter provider
return MeterProvider(resource=resource, metric_readers=[prometheus_reader, periodic_reader])
except Exception as e:
logger.error(f"Failed to set up metrics provider: {str(e)}")
# Return a no-op meter provider
return MeterProvider(resource=resource)
def initialize_telemetry(
tracing_config: Optional[TracingConfig] = None,
metrics_config: Optional[MetricsConfig] = None,
) -> Tracer:
"""
Initialize OpenTelemetry with tracing and metrics.
Args:
tracing_config: Configuration for tracing
metrics_config: Configuration for metrics
Returns:
Tracer: Configured OpenTelemetry tracer
Example:
```python
config = TracingConfig(
service_name="amq-adapter",
service_version="1.0.0",
otlp_endpoint="http://jaeger:4317",
debug=True
)
metrics_config = MetricsConfig(
prometheus_port=9464,
prometheus_host="0.0.0.0"
)
tracer = initialize_telemetry(config, metrics_config)
```
"""
# Use default configs if none provided
if tracing_config is None:
tracing_config = TracingConfig(
service_name=os.getenv("SERVICE_NAME", "amq-adapter"),
service_version=os.getenv("SERVICE_VERSION", "unknown"),
otlp_endpoint=os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317"),
insecure=os.getenv("OTEL_INSECURE_MODE", "true").lower() == "true",
debug=os.getenv("OTEL_DEBUG", "false").lower() == "true",
)
if metrics_config is None:
metrics_config = MetricsConfig(
prometheus_port=int(os.getenv("PROMETHEUS_PORT", "9464")),
prometheus_host=os.getenv("PROMETHEUS_HOST", "localhost"),
export_interval_millis=int(os.getenv("METRICS_EXPORT_INTERVAL", "30000")),
)
try:
# Create resource
resource = create_resource(tracing_config)
# Set up tracing
provider = setup_trace_provider(tracing_config, resource)
trace.set_tracer_provider(provider)
logger.info(f"Initialized tracing for service {tracing_config.service_name}")
# Set up metrics
metrics_provider = setup_metrics_provider(metrics_config, resource)
metrics.set_meter_provider(metrics_provider)
logger.info("Initialized metrics export")
# Return configured tracer
return trace.get_tracer(
tracing_config.service_name,
tracing_config.service_version,
)
except Exception as e:
logger.error(f"Failed to initialize telemetry: {str(e)}")
# Return a no-op tracer in case of failure
return trace.get_tracer("no-op")
def create_span(tracer: Tracer, name: str, attributes: Optional[Dict[str, Any]] = None):
"""
Create a new span with the given name and attributes.
Args:
tracer: The OpenTelemetry tracer
name: Name of the span
attributes: Optional attributes to add to the span
Returns:
A context manager that creates and manages a span
"""
return tracer.start_as_current_span(
name,
attributes=attributes or {},
)
+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
self.amq_service: AMQService = AMQService(amq_configuration, self)
self.thread = Thread(target=self.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
+36 -259
View File
@@ -5,300 +5,77 @@ It is NOT MEANT to be used as a standalone module. It WILL NOT run outside the C
""" """
import importlib import importlib
import pathlib
import re
from collections import defaultdict
from threading import Thread from threading import Thread
from typing import Optional, List, Dict from typing import Any, List, Optional
from aiohttp import ClientSession
from fastapi.routing import APIRoute
from fastapi.security import OAuth2PasswordRequestForm
from starlette.datastructures import UploadFile, Headers
# CleverSwarm specific imports # CleverSwarm specific imports
import core.deps import core.deps
from rest.v0 import api from fastapi.routing import APIRoute
from schemas.user_schema import UserSchema from schemas.user_schema import UserSchema
from amqp.adapter.data_parser import AMQDataParser from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter
# AMQP specific imports
from amqp.adapter.data_response_factory import AMQResponseFactory
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter, AMQMessage
from amqp.config.amq_configuration import AMQConfiguration from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQResponse, AMQErrorMessage
from amqp.service.amq_service import AMQService 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" 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. Get the function/Callable for the module_name.function_name if it exists in the module, else return None.
It also adjusts the path regex to ensure it matches the path in different API versions.
""" """
# Check if the route has a preferred endpoint with json serialization
_preferred_endpoint = get_prefered_endpoint(
getattr(original_route.endpoint, "__module__"),
original_route.endpoint.__name__ + PREFERRED_SUFFIX,
)
# Need to adapt also the path regex to ensure it matches the path in different API versions
# Ensures the leading '^' is followed by any match and there's no '/' before final '$'
_new_regex_string = (
"^.*" + original_route.path_regex.pattern.lstrip("^").rstrip("/$") + "$"
)
_new_route: APIRoute = APIRoute(
path=original_route.path,
endpoint=(
_preferred_endpoint if _preferred_endpoint else original_route.endpoint
),
response_model=original_route.response_model,
status_code=original_route.status_code,
tags=original_route.tags.copy() if original_route.tags else None,
dependencies=(
original_route.dependencies.copy() if original_route.dependencies else None
),
summary=original_route.summary,
description=original_route.description,
response_description=original_route.response_description,
responses=original_route.responses.copy() if original_route.responses else None,
deprecated=original_route.deprecated,
name=original_route.name,
methods=original_route.methods.copy() if original_route.methods else None,
operation_id=original_route.operation_id,
response_model_include=original_route.response_model_include,
response_model_exclude=original_route.response_model_exclude,
response_model_by_alias=original_route.response_model_by_alias,
response_model_exclude_unset=original_route.response_model_exclude_unset,
response_model_exclude_defaults=original_route.response_model_exclude_defaults,
response_model_exclude_none=original_route.response_model_exclude_none,
include_in_schema=original_route.include_in_schema,
response_class=original_route.response_class,
dependency_overrides_provider=original_route.dependency_overrides_provider,
callbacks=original_route.callbacks.copy() if original_route.callbacks else None,
openapi_extra=(
original_route.openapi_extra.copy()
if original_route.openapi_extra
else None
),
generate_unique_id_function=original_route.generate_unique_id_function,
)
_new_route.path_regex = re.compile(_new_regex_string)
return _new_route
def get_prefered_endpoint(module_name: str, function_name: str) -> Optional[callable]:
"""Get the function if it exists in the module, else return None."""
try: try:
_module = importlib.import_module(module_name) _module = importlib.import_module(module_name)
_prefered_endpoint = getattr(_module, function_name) _preferred_endpoint = getattr(_module, function_name)
return _prefered_endpoint if callable(_prefered_endpoint) else None return _preferred_endpoint if callable(_preferred_endpoint) else None
except Exception: except Exception:
return None return None
def group_routes_by_method(api_routes: List[APIRoute]) -> Dict[str, List[APIRoute]]: # =================================================================================================
""" # ======================== C l e v e r S w a r m A M Q A d a p t e r ========================
Groups API routes by their HTTP method. This is important to ensure that the URLPath is matched # =================================================================================================
correctly, as different methods may share the same URLPath (e.g. GET, POST PUT, DELETE).
"""
_method_to_routes = defaultdict(
list
) # Automatically initializes new lists for new keys
for _route in api_routes:
if (
_route.methods
): # Ensure methods exist (should always be true for valid routes)
for method in _route.methods:
_method_to_routes[method].append(clone_and_adapt_route(_route))
return dict(_method_to_routes)
class CleverSwarmAmqpAdapter(CleverThisServiceAdapter): 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__( def __init__(self, amq_configuration: AMQConfiguration, routes: List[APIRoute]):
self, amq_configuration: AMQConfiguration, session: ClientSession = None super().__init__(amq_configuration, routes=routes, session=None)
): self.amq_service: AMQService = AMQService(amq_configuration, self)
super().__init__(amq_configuration, session) self.thread = Thread(target=self.amq_service.run)
self.routes = group_routes_by_method(api_routes=api.api_router.routes)
_amq_service: AMQService = AMQService(amq_configuration, self)
self.thread = Thread(target=_amq_service.run)
self.thread.start() self.thread.start()
async def get_current_user(self, token: str) -> UserSchema: async def get_current_user(self, token: str) -> UserSchema:
""" """
Overrides the default get_current_user method to use the token from the AMQP message. 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 :param token: JWT token from the AMQP message
:return: UserSchema object :return: UserSchema object
""" """
return await core.deps.get_current_user(token) return await core.deps.get_current_user(token)
# ================================================================================================== def get_endpoint(self, endpoint: Any) -> Any:
# ===================== C l e v e r S w a r m A P I m e t h o d s ============================ """
# ================================================================================================== 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: :param endpoint: the Callable for the endpoint
"""
Handles the GET requests for the CleverSwarm API.
:param message: message from the AMQP that contains the GET request
:param auth_user: user authenticated by the API Gateway, as UserSchem
:return: AMQResponse to be sent back to the API Gateway via AMQP
"""
return await self.handle_no_body_payload(message, auth_user, "GET")
async def put(self, message: AMQMessage, auth_user: UserSchema) -> AMQResponse: :return: new Callable representing the JSON wrapper endpoint, if it exists or the original Callable.
""" """
Handles the PUT requests for the CleverSwarm API. _preferred_endpoint = get_callable_for_name(
:param message: message from the AMQP that contains the PUT request getattr(endpoint, "__module__"),
:param auth_user: user authenticated by the API Gateway, as UserSchem endpoint.__name__ + PREFERRED_SUFFIX,
: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_error_message(
message.id,
AMQErrorMessage(
404,
"Not Found",
f"Destination location [{message.path}] does not exists",
),
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_error_message(
message.id,
AMQErrorMessage(
404,
"Not Found",
f"Destination location [{message.path}] does not exists",
),
message.trace_info,
) )
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 clever_swarm_adapter import CleverSwarmAmqpAdapter
from amqp.config.amq_configuration import AMQConfiguration
if __name__ == "__main__": if __name__ == "__main__":
adapter = CleverSwarmAmqpAdapter(AMQConfiguration("./application.properties.local")) adapter = CleverSwarmAmqpAdapter(AMQConfiguration("./application.properties.local"))
print("Press ^C to exit...") print("Press ^C to exit...")
+353
View File
@@ -0,0 +1,353 @@
import json
import logging
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/v4/documents")
async def get_documents():
# Create a custom span for the logic within this endpoint
documents = []
try:
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)):
did = f"doc-{random.randint(1000, 9999)}"
documents.append(
{
"id": did,
"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)),
),
}
)
otadapter.amq_service.command(
"otdemo-2",
"print_document_id_static",
document_id=did + "-static",
)
otadapter.amq_service.command(
"otdemo-2",
"print_document_id",
document_id=did,
)
# Increment the documents_retrieved_total metric
documents_retrieved_counter.add(1)
except Exception as e:
logging.error(f"Error retrieving documents: {e}")
# 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:
+21
View File
@@ -0,0 +1,21 @@
# ====================================================================
# CleverMicro AMQ Adapter settings
# ====================================================================
[CleverMicro-AMQ]
cm.amq-adapter.service-name=otdemo-2
cm.amq-adapter.generator-id=1
cm.amq-adapter.route-mapping=otdemo-2:otdemo-exchange:otdemo-queue-2:otdemo-2.#:5:0
cm.amq-adapter.require-authenticated-user=false
cm.amq-adapter.callbacks-enabled=False
# CleverMicro AMQ Adapter dispatcher, settings related to RabbitMQ connection and optional REST forwarding
cm.dispatch.use-dlq=true
cm.dispatch.amq-host=localhost
cm.dispatch.amq-port=5672
cm.dispatch.amq-port-tls=5671
cm.dispatch.download-dir=/tmp/downloaded_files
cm.backpressure.threshold=1
cm.backpressure.time-window=150
cm.backpressure.cpu-overload-duration=300
cm.backpressure.cpu-idle-duration=300
+21
View File
@@ -0,0 +1,21 @@
# ====================================================================
# CleverMicro AMQ Adapter settings
# ====================================================================
[CleverMicro-AMQ]
cm.amq-adapter.service-name=otdemo
cm.amq-adapter.generator-id=1
cm.amq-adapter.route-mapping=otdemo:otdemo-exchange:otdemo-queue-1:api.cm.dev.cleverthis.localhost.test.python.frontend.#:5:0
cm.amq-adapter.require-authenticated-user=false
cm.amq-adapter.callbacks-enabled=True
# CleverMicro AMQ Adapter dispatcher, settings related to RabbitMQ connection and optional REST forwarding
cm.dispatch.use-dlq=true
cm.dispatch.amq-host=localhost
cm.dispatch.amq-port=5672
cm.dispatch.amq-port-tls=5671
cm.dispatch.download-dir=/tmp/downloaded_files
cm.backpressure.threshold=1
cm.backpressure.time-window=150
cm.backpressure.cpu-overload-duration=300
cm.backpressure.cpu-idle-duration=300
+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)
self.amq_service: AMQService = AMQService(amq_configuration, self)
self.thread = Thread(target=self.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
+53
View File
@@ -0,0 +1,53 @@
import asyncio
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQRoute, DataMessage
from otdemo.otdemo_adapter import OTDemoAmqpAdapter
def print_document_id_static(document_id: str) -> str:
print(f"STATIC Document ID: {document_id}")
return "STATIC PRINTED!"
class Backend:
def print_document_id(self, document_id: str) -> str:
print(f"CLASS Document ID: {document_id}")
return "CLASS PRINTED!"
async def processing_loop():
while (
not otadapter.amq_service.future.done()
): # Wait until AMQ adapter initializes RabbitMQ connection
await asyncio.sleep(0.1)
# IMPORTANT - in the next line the string arg must match the first token in route-mapping property
route: AMQRoute = otadapter.amq_service.router.find_route_by_service("otdemo-2")
generator = otadapter.amq_service.data_message_generator(route) # this is a generator function
print(f"Data message generator created: {generator}")
message = await generator.asend(None) # First value must be None to start the generator
while True:
print(f"Received message: {message}")
if message is None:
break
data_message: DataMessage = message
# This is a message processing part. You may use data_message.method() or .domain() or .path()
# to access values provided by caller in the send command() call.
document_id = data_message._base64body.decode("utf-8")
print(
f"DATA_MESSAGE Attributes: {data_message.method()} {data_message.domain()} {data_message.path()}"
)
if data_message.path() == "print_document_id_static":
print_document_id_static(document_id)
elif data_message.path() == "print_document_id":
backend = Backend()
backend.print_document_id(document_id)
message = await generator.asend(True)
if __name__ == "__main__":
otadapter = OTDemoAmqpAdapter(AMQConfiguration("otdemo.backend.properties"), [])
# Run the processing loop in the AMQ service event loop
asyncio.run_coroutine_threadsafe(processing_loop(), otadapter.amq_service.loop)
otadapter.thread.join() # if the code above is not running a loop, need to join the thread to keep program running
+16 -4
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project] [project]
name = "amq_adapter" name = "amq_adapter"
version = "0.2.30" version = "0.2.37"
description = "The Python implementation of the AMQ Adapter for CleverMicro" description = "The Python implementation of the AMQ Adapter for CleverMicro"
readme = "README.md" readme = "README.md"
authors = [ authors = [
@@ -29,18 +29,30 @@ dependencies = [
"opentelemetry-instrumentation-pika", "opentelemetry-instrumentation-pika",
"pika", "pika",
"psutil", "psutil",
"python_multipart" "python-multipart",
"consul-kv"
] ]
[tool.setuptools.packages.find] [tool.setuptools.packages.find]
where = ["."] where = ["."]
include = ["amqp*"] include = ["amqp*"]
[tool.black]
line-length = 100
[tool.pytest.ini_options]
asyncio_default_fixture_loop_scope = "module"
[project.optional-dependencies] [project.optional-dependencies]
dev = [ dev = [
"pytest>=7.0", "pytest>=7.0",
"pytest-cov>=4.0", "pytest-cov>=4.0",
"pytest-asyncio",
"build", "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 import unittest
from enum import Enum
from typing import Any, List, Optional
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import AsyncMock, MagicMock, patch
from aio_pika.abc import AbstractChannel, AbstractRobustConnection
from aiohttp import ClientResponse 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.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): def setUp(self):
# Mock AMQConfiguration # Mock AMQConfiguration
self.mock_amq_config = AMQConfiguration("./application.properties.local") self.mock_amq_config = AMQConfiguration("./application.properties.local")
@@ -18,11 +109,13 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
# Create adapter instance # Create adapter instance
self.adapter = CleverThisServiceAdapter( 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): 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.assertEqual(self.adapter.service_port, 8080)
self.assertFalse(self.adapter.require_authenticated_user) self.assertFalse(self.adapter.require_authenticated_user)
@@ -48,26 +141,27 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
@patch.object(CleverThisServiceAdapter, "get_current_user") @patch.object(CleverThisServiceAdapter, "get_current_user")
async def test_on_message_unauthenticated(self, mock_get_user): async def test_on_message_unauthenticated(self, mock_get_user):
mock_message = MagicMock(spec=AMQMessage) mock_message = MagicMock(spec=AMQMessage)
mock_message.headers = {} mock_message.headers.return_value = {}
mock_message.path = "/some/path" mock_message.path.return_value = "/some/path"
mock_message.method = "GET" mock_message.method.return_value = "GET"
mock_message.id = "123" mock_message.id.return_value = "123"
mock_message.trace_info = {} mock_message.trace_info.return_value = {}
# Test when authentication is required but not provided # Test when authentication is required but not provided
self.adapter.require_authenticated_user = True
response = await self.adapter.on_message(mock_message) 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() mock_get_user.assert_not_called()
@patch.object(CleverThisServiceAdapter, "get_current_user") @patch.object(CleverThisServiceAdapter, "get_current_user")
@patch.object(CleverThisServiceAdapter, "get") @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 = MagicMock(spec=AMQMessage)
mock_message.headers = {"Authorization": ["Bearer valid_token"]} mock_message.headers.return_value = {"Authorization": ["Bearer valid_token"]}
mock_message.path = "/some/path" mock_message.path.return_value = "/some/path"
mock_message.method = "GET" mock_message.method.return_value = "GET"
mock_message.id = "123" mock_message.id.return_value = "123"
mock_message.trace_info = {} mock_message.trace_info.return_value = {}
mock_user = MagicMock() mock_user = MagicMock()
mock_get_user.return_value = mock_user mock_get_user.return_value = mock_user
@@ -75,6 +169,7 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
mock_response = MagicMock(spec=DataResponse) mock_response = MagicMock(spec=DataResponse)
mock_get.return_value = mock_response mock_get.return_value = mock_response
self.adapter.require_authenticated_user = True
response = await self.adapter.on_message(mock_message) response = await self.adapter.on_message(mock_message)
self.assertEqual(response, mock_response) self.assertEqual(response, mock_response)
@@ -84,16 +179,14 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
@patch.object(CleverThisServiceAdapter, "get_current_user") @patch.object(CleverThisServiceAdapter, "get_current_user")
async def test_on_message_login_path(self, mock_get_user): async def test_on_message_login_path(self, mock_get_user):
mock_message = MagicMock(spec=AMQMessage) mock_message = MagicMock(spec=AMQMessage)
mock_message.headers = {} mock_message.headers.return_value = {}
mock_message.path = "/login" mock_message.path.return_value = "/login"
mock_message.method = "POST" mock_message.method.return_value = "POST"
mock_message.id = "123" mock_message.id.return_value = "123"
mock_message.trace_info = {} mock_message.trace_info.return_value = {}
mock_response = MagicMock(spec=DataResponse) mock_response = MagicMock(spec=DataResponse)
with patch.object( with patch.object(self.adapter, "post", return_value=mock_response) as mock_post:
self.adapter, "post", return_value=mock_response
) as mock_post:
response = await self.adapter.on_message(mock_message) response = await self.adapter.on_message(mock_message)
self.assertEqual(response, mock_response) self.assertEqual(response, mock_response)
mock_get_user.assert_not_called() mock_get_user.assert_not_called()
@@ -110,11 +203,11 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
for method in methods: for method in methods:
with self.subTest(method=method): with self.subTest(method=method):
mock_message = MagicMock(spec=AMQMessage) mock_message = MagicMock(spec=AMQMessage)
mock_message.headers = {"Authorization": ["Bearer valid_token"]} mock_message.headers.return_value = {"Authorization": ["Bearer valid_token"]}
mock_message.path = "/some/path" mock_message.path.return_value = "/some/path"
mock_message.method = method mock_message.method.return_value = method
mock_message.id = "123" mock_message.id.return_value = "123"
mock_message.trace_info = {} mock_message.trace_info.return_value = {}
with patch.object( with patch.object(
self.adapter, method.lower(), return_value=mock_response self.adapter, method.lower(), return_value=mock_response
@@ -124,11 +217,11 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
mock_method.assert_called_once() mock_method.assert_called_once()
async def test_get_with_session(self): async def test_get_with_session(self):
mock_message = MagicMock(spec=DataMessage) mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.path = "/api/test" mock_message.path.return_value = "/api/test"
mock_message.headers = {"X-Test": "value"} mock_message.headers.return_value = {"X-Test": "value"}
mock_message.trace_info = {"trace-id": "123"} mock_message.trace_info.return_value = {"trace-id": "123"}
mock_message.id = "msg123" mock_message.id.return_value = "msg123"
mock_response = MagicMock(spec=ClientResponse) mock_response = MagicMock(spec=ClientResponse)
mock_response.status = 200 mock_response.status = 200
@@ -139,7 +232,7 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
response = await self.adapter.get(mock_message, None) 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( self.mock_session.get.assert_called_once_with(
"http://testhost:8080/api/test", "http://testhost:8080/api/test",
headers={"X-Test": "value", "trace-id": "123"}, headers={"X-Test": "value", "trace-id": "123"},
@@ -148,50 +241,50 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
async def test_get_without_session(self): async def test_get_without_session(self):
self.adapter.session = None self.adapter.session = None
mock_message = MagicMock(spec=DataMessage) mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.path = "/api/test" mock_message.path.return_value = "/api/test"
mock_message.id = "msg123" mock_message.id.return_value = "msg123"
mock_message.trace_info = {} mock_message.trace_info.return_value = {}
response = await self.adapter.get(mock_message, None) 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): async def test_put_with_form_handling(self):
mock_message = MagicMock(spec=DataMessage) mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.path = "/api/test" mock_message.path.return_value = "/api/test"
mock_message.headers = {} mock_message.headers.return_value = {}
mock_message.trace_info = {} mock_message.trace_info.return_value = {}
mock_message.id = "msg123" mock_message.id.return_value = "msg123"
mock_message.body = b"test=value" mock_message.body.return_value = b"test=value"
mock_response = MagicMock(spec=ClientResponse) mock_response: DataResponse = MagicMock(spec=DataResponse)
mock_response.status = 200 mock_response.response_code.return_value = "200"
mock_response.content_type = "application/json" mock_response.content_type.return_value = "application/json"
mock_response.read = AsyncMock(return_value=b'{"result": "success"}') mock_response.read = AsyncMock(return_value=b'{"result": "success"}')
self.mock_session.put.return_value = mock_response self.mock_session.put.return_value = mock_response
with patch.object( with patch.object(
self.adapter, self.adapter,
"handle_possible_form", "_authorize_forward_request",
wraps=self.adapter.handle_possible_form, wraps=self.adapter._authorize_forward_request,
) as mock_handler: ) as mock_handler:
response = await self.adapter.put(mock_message, None) 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() mock_handler.assert_called_once()
async def test_post_with_form_handling(self): async def test_post_with_form_handling(self):
mock_message = MagicMock(spec=DataMessage) mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.path = "/api/test" mock_message.path.return_value = "/api/test"
mock_message.headers = {} mock_message.headers.return_value = {}
mock_message.trace_info = {} mock_message.trace_info.return_value = {}
mock_message.id = "msg123" mock_message.id.return_value = "msg123"
mock_message.body = b"test=value" mock_message.body.return_value = b"test=value"
mock_response = MagicMock(spec=ClientResponse) mock_response: DataResponse = MagicMock(spec=DataResponse)
mock_response.status = 201 mock_response.response_code.return_value = "201"
mock_response.content_type = "application/json" mock_response.content_type = "application/json"
mock_response.read = AsyncMock(return_value=b'{"result": "created"}') mock_response.read = AsyncMock(return_value=b'{"result": "created"}')
@@ -199,28 +292,139 @@ class TestCleverThisServiceAdapter(unittest.TestCase):
with patch.object( with patch.object(
self.adapter, self.adapter,
"handle_possible_form", "_authorize_forward_request",
wraps=self.adapter.handle_possible_form, wraps=self.adapter._authorize_forward_request,
) as mock_handler: ) as mock_handler:
response = await self.adapter.post(mock_message, None) 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() mock_handler.assert_called_once()
async def test_head_not_implemented(self): async def test_head_not_implemented(self):
mock_message = MagicMock(spec=DataMessage) mock_message: AMQMessage = MagicMock(spec=AMQMessage)
mock_message.id = "msg123" mock_message.id = "msg123"
with self.assertRaises(NotImplementedError): with self.assertRaises(NotImplementedError):
await self.adapter.head(mock_message, None) await self.adapter.head(mock_message, None)
async def test_delete_not_implemented(self): def test_basic_type_conversion(self):
mock_message = MagicMock(spec=DataMessage) async def test_method(x: int, y: float, z: bool, text: str):
mock_message.id = "msg123" pass
# Assuming delete is not implemented in base class args = {"x": "123", "y": "45.67", "z": "true", "text": "hello"}
with self.assertRaises(NotImplementedError):
await self.adapter.delete(mock_message, None) 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__": 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")) msg = DataMessageFactory.from_bytes(raw.encode("utf-8"))
self.assertIsNotNone(msg.id) self.assertIsNotNone(msg.id)
self.assertEqual( self.assertEqual(
msg.id.as_string(), msg.id().as_string(),
"64F3B5AF4D00000A4000".lower(), "64F3B5AF4D00000A4000".lower(),
"SnowflakeID incorrectly parsed.", "SnowflakeID incorrectly parsed.",
) )
self.assertEqual( self.assertEqual(
msg.headers["Content-Type"][0], msg.headers()["Content-Type"][0],
"multipart/form-data; boundary=------------------------j1RjIqLrFH07XMsyMvpS1v", "multipart/form-data; boundary=------------------------j1RjIqLrFH07XMsyMvpS1v",
) )
def test_generate_snowflake_id(self): def test_generate_snowflake_id(self):
_id = DataMessageFactory.get_instance(1).generate_snowflake_id() _id = DataMessageFactory.get_instance(1).generate_snowflake_id()
self.assertEqual( self.assertEqual("00001000", _id.as_string()[-8:], "Instance ID not set correctly")
"00001000", _id.as_string()[-8:], "Instance ID not set correctly"
)
_id = DataMessageFactory(2).generate_snowflake_id() _id = DataMessageFactory(2).generate_snowflake_id()
self.assertEqual( self.assertEqual("00002000", _id.as_string()[-8:], "Instance ID not set correctly")
"00002000", _id.as_string()[-8:], "Instance ID not set correctly"
)
_f = DataMessageFactory(127) _f = DataMessageFactory(127)
_i0 = _f.generate_snowflake_id() _i0 = _f.generate_snowflake_id()
@@ -65,9 +61,9 @@ class DataMessageFactoryTest(TestCase):
{"k1": ["v1"], "k2": ["v2"]}, {"k1": ["v1"], "k2": ["v2"]},
"Quick brown fox", "Quick brown fox",
) )
self.assertEqual("localhost", _req.domain, "Domain 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("/test/me", _req.path(), "Context path not set correctly")
self.assertEqual("v1", _req.headers["k1"][0], "Headers 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") self.assertEqual(b"Quick brown fox", _req.body(), "Body not correctly decoded")
def test_serialize(self): def test_serialize(self):
@@ -82,11 +78,11 @@ class DataMessageFactoryTest(TestCase):
_ser = _f.serialize(_req) _ser = _f.serialize(_req)
self.assertEqual(65, _ser[2]) self.assertEqual(65, _ser[2])
_deser = _f.from_bytes(_ser) _deser = _f.from_bytes(_ser)
self.assertEqual(_req.id.as_string(), _deser.id.as_string()) self.assertEqual(_req.id().as_string(), _deser.id().as_string())
self.assertEqual("localhost", _deser.domain) self.assertEqual("localhost", _deser.domain())
self.assertEqual("/test/me", _deser.path) self.assertEqual("/test/me", _deser.path())
self.assertEqual(b"Quick brown fox", _deser.body()) 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): def test_amq_message_routing_key(self):
_f = DataMessageFactory(35) _f = DataMessageFactory(35)
@@ -111,7 +107,7 @@ class DataMessageFactoryTest(TestCase):
self.assertLessEqual( self.assertLessEqual(
_x_timestamp - _timestamp, _x_timestamp - _timestamp,
1, 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): def test_to_string(self):
@@ -124,6 +120,4 @@ class DataMessageFactoryTest(TestCase):
"Quick brown fox", "Quick brown fox",
) )
_as_str = _f.to_string(_req) _as_str = _f.to_string(_req)
self.assertGreater( self.assertGreater(_as_str.index("domain='localhost'"), 0, "Expected content not found")
_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_message_factory import DataMessageFactory
from amqp.adapter.data_response_factory import DataResponseFactory from amqp.adapter.data_response_factory import DataResponseFactory
from amqp.model.model import DataMessageMagicByte
from amqp.model.snowflake_id import SnowflakeId from amqp.model.snowflake_id import SnowflakeId
@@ -11,8 +12,8 @@ class TestDataResponseFactory(TestCase):
_f = DataResponseFactory() _f = DataResponseFactory()
_g = DataMessageFactory(111) _g = DataMessageFactory(111)
_resp = _f.create_async_response_message(_g.generate_snowflake_id()) _resp = _f.create_async_response_message(_g.generate_snowflake_id())
self.assertEqual(200, _resp.response_code) self.assertEqual("200", _resp.response_code())
self.assertEqual(b"OK", _resp.response) self.assertEqual(b"OK", _resp.body())
def test_serialize(self): def test_serialize(self):
_f = DataResponseFactory() _f = DataResponseFactory()
@@ -22,7 +23,7 @@ class TestDataResponseFactory(TestCase):
_ser = _f.serialize(_resp) _ser = _f.serialize(_resp)
self.assertEqual( self.assertEqual(
2, 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", "Magic byte and SnowflakeId not found",
) )
self.assertGreater(str(_ser[2:]).index("|e=|"), 2) self.assertGreater(str(_ser[2:]).index("|e=|"), 2)
@@ -35,6 +36,6 @@ class TestDataResponseFactory(TestCase):
_resp = _f.create_async_response_message(_snowflakeId) _resp = _f.create_async_response_message(_snowflakeId)
_ser = _f.serialize(_resp) _ser = _f.serialize(_resp)
_deser = _f.from_bytes(_ser) _deser = _f.from_bytes(_ser)
self.assertEqual(_snowflakeId.as_string(), _deser.id) self.assertEqual(_snowflakeId.as_string(), _deser.id())
self.assertEqual(200, _deser.response_code) self.assertEqual("200", _deser.response_code())
self.assertEqual(b"OK", _deser.response) self.assertEqual(b"OK", _deser.body())
+263 -2
View File
@@ -1,11 +1,22 @@
import asyncio
import os import os
import tempfile 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 from amqp.adapter.file_handler import StreamingFileHandler
class TestFileHandler(TestCase): class TestFileHandler(IsolatedAsyncioTestCase):
def test_ensure_directory_exists(self): def test_ensure_directory_exists(self):
""" """
Test function for ensure_directory_exists. Creates a temporary directory 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')}" ), f"Test Case 3 Version 3 Failed: {result3_v3} != {os.path.join(temp_dir, 'test3.3.txt')}"
print("All test cases passed!") 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 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): 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): 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): def test_deserialize_object(self):
assert True serialized = serialize_object(self.test_model)
deserialized = deserialize_object(serialized)
def test__convert_value(self): self.assertIsInstance(deserialized, TestModel)
assert True 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): def test_serialize_deserialize_with_none(self):
assert True 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): def test_serialize_deserialize_with_list_of_models(self):
assert True 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): def test_parse_url_query_no_params(self):
assert True test_url = "http://example.com/path"
path, params = parse_url_query(test_url)
def test_serialize_to_json(self): self.assertEqual(path, "/path")
assert True 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): def test_parse_url_query(self):
path, args = parse_url_query( path, args = parse_url_query(
"http://localhost:8080/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741?response_type=JsonFile" "http://localhost:8080/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741?response_type=JsonFile"
) )
assert path == "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741" self.assertEqual(path, "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741")
assert args == {"response_type": "JsonFile"} self.assertEqual(args, {"response_type": "JsonFile"})
path, args = parse_url_query( path, args = parse_url_query(
"http://localhost:8080/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741" "http://localhost:8080/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741"
) )
assert path == "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741" self.assertEqual(path, "/api/v0/unstructured/863e4d9f-4612-4173-8d0a-3de34cc7a741")
assert args == {} 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
+145
View File
@@ -0,0 +1,145 @@
import importlib
import sys
import unittest
from amqp.service.amq_message_handler import invoke_command
class TestRPCInvocation(unittest.IsolatedAsyncioTestCase):
"""Test cases for the RPC invocation functionality in DataMessageHandler."""
async def asyncSetUp(self):
"""Set up the test environment with a dummy module."""
# Define a simple module string
dummy_module_content = """
import asyncio
# Global function
def greet_sync(name: str) -> str:
return f"Hello, {name} (sync global)!"
async def greet_async(name: str) -> str:
await asyncio.sleep(0.01) # Simulate async work
return f"Hello, {name} (async global)!"
# A simple class
class MyService:
def __init__(self):
self.service_name = "MyAwesomeService"
def process_data_sync(self, data: str, count: int) -> str:
return f"Service '{self.service_name}' processed '{data}' {count} times (sync method)."
async def process_data_async(self, data: str, delay: float) -> str:
await asyncio.sleep(delay) # Simulate async work
return f"Service '{self.service_name}' processed '{data}' after {delay}s (async method)."
def _private_method(self):
return "This is a private method."
"""
# Dynamically create a module object
spec = importlib.util.spec_from_loader("my_package.my_module", loader=None)
if spec is None:
raise ImportError("Could not create module spec for my_package.my_module")
self.my_module = importlib.util.module_from_spec(spec)
exec(dummy_module_content, self.my_module.__dict__)
sys.modules["my_package.my_module"] = (
self.my_module
) # Add to sys.modules for importlib to find
async def asyncTearDown(self):
"""Clean up after tests."""
# Remove the dummy module from sys.modules
if "my_package" in sys.modules:
del sys.modules["my_package"]
async def test_invoke_global_sync_function(self):
"""Test invoking a global synchronous function."""
result = await invoke_command(
package_name="my_package.my_module",
class_name="__global__",
function_name="greet_sync",
kwargs={"name": "Alice"},
)
self.assertEqual(result.body().decode(), "Hello, Alice (sync global)!")
async def test_invoke_global_async_function(self):
"""Test invoking a global asynchronous function."""
result = await invoke_command(
package_name="my_package.my_module",
class_name="__global__",
function_name="greet_async",
kwargs={"name": "Bob"},
)
self.assertEqual(result.body().decode(), "Hello, Bob (async global)!")
async def test_invoke_class_sync_method(self):
"""Test invoking a class synchronous method."""
result = await invoke_command(
package_name="my_package.my_module",
class_name="MyService",
function_name="process_data_sync",
kwargs={"data": "report", "count": 3},
)
self.assertEqual(
result.body().decode(),
"Service 'MyAwesomeService' processed 'report' 3 times (sync method).",
)
async def test_invoke_class_async_method(self):
"""Test invoking a class asynchronous method."""
result = await invoke_command(
package_name="my_package.my_module",
class_name="MyService",
function_name="process_data_async",
kwargs={"data": "metrics", "delay": 0.01},
)
self.assertEqual(
result.body().decode(),
"Service 'MyAwesomeService' processed 'metrics' after 0.01s (async method).",
)
async def test_package_not_found(self):
"""Test error handling when package is not found."""
with self.assertRaises(ModuleNotFoundError):
await invoke_command(
package_name="non_existent_package",
class_name="__global__",
function_name="some_func",
kwargs={},
)
async def test_class_not_found(self):
"""Test error handling when class is not found."""
with self.assertRaises(AttributeError):
await invoke_command(
package_name="my_package.my_module",
class_name="NonExistentClass",
function_name="some_method",
kwargs={},
)
async def test_method_not_found(self):
"""Test error handling when method is not found."""
with self.assertRaises(AttributeError):
await invoke_command(
package_name="my_package.my_module",
class_name="MyService",
function_name="non_existent_method",
kwargs={},
)
async def test_private_method_access(self):
"""Test that private methods can be accessed (Python doesn't enforce privacy)."""
result = await invoke_command(
package_name="my_package.my_module",
class_name="MyService",
function_name="_private_method",
kwargs={},
)
self.assertEqual(result.body().decode(), "This is a private method.")
if __name__ == "__main__":
unittest.main()
+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 unittest import TestCase
from amqp.adapter.service_message_factory import ( from amqp.adapter.service_message_factory import (
ServiceMessageFactory,
RS,
INVALID_MESSAGE, INVALID_MESSAGE,
RS,
ServiceMessageFactory,
) )
from amqp.model.model import ServiceMessage from amqp.model.model import ServiceMessage
from amqp.model.service_message_type import ServiceMessageType from amqp.model.service_message_type import ServiceMessageType
@@ -13,15 +13,11 @@ class ServiceMessageTest(TestCase):
def test_to_string(self): def test_to_string(self):
_f: ServiceMessageFactory = ServiceMessageFactory("bumble-bee") _f: ServiceMessageFactory = ServiceMessageFactory("bumble-bee")
message: ServiceMessage = _f.of( message: ServiceMessage = _f.of(ServiceMessageType.ROUTING_DATA, "Duck", "Donald")
ServiceMessageType.ROUTING_DATA, "Duck", "Donald"
)
message_str = _f.to_string(message) message_str = _f.to_string(message)
self.assertTrue("|type=01|" in message_str) self.assertTrue("|type=01|" in message_str)
self.assertTrue("|body=Duck" in message_str) self.assertTrue("|body=Duck" in message_str)
message: ServiceMessage = _f.of( message: ServiceMessage = _f.of(ServiceMessageType.ROUTING_DATA_REPLACE, "Duck", "Donald")
ServiceMessageType.ROUTING_DATA_REPLACE, "Duck", "Donald"
)
message_str = _f.to_string(message) message_str = _f.to_string(message)
self.assertTrue("|type=03|" in message_str) self.assertTrue("|type=03|" in message_str)
self.assertTrue("|body=Duck" 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") message = _f.of(ServiceMessageType.ROUTING_DATA, "Duck", "Donald")
serialized = _f.serialize(message) serialized = _f.serialize(message)
assert RS + "0" + str(ServiceMessageType.ROUTING_DATA.value) 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 import pytest
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.amq_route_factory import AMQRouteFactory from amqp.adapter.amq_route_factory import AMQRouteFactory
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.model.model import AMQRoute from amqp.model.model import AMQRoute
from amqp.router.route_database import RouteDatabase from amqp.router.route_database import RouteDatabase
@@ -80,7 +80,7 @@ class TestRouteDatabase(TestCase):
message = DataMessageFactory.get_instance(111).create_request_message( message = DataMessageFactory.get_instance(111).create_request_message(
"POST", "clever3-book.localhost", "/aa/bb?cc=d", {}, "" "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.assertIsNotNone(route)
self.assertEqual(ROUTE_3, route.as_string()) self.assertEqual(ROUTE_3, route.as_string())
route = self._database.find_route("clever3.book.localhost", "/three-hot-dogs") 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()