From 5ab9406430d6de37a35fcb8c6a23144e4634891a Mon Sep 17 00:00:00 2001 From: Stanislav Hejny Date: Fri, 18 Jul 2025 07:47:21 +0100 Subject: [PATCH] feat: implement UserManagementServiceClient with RabbitMQ communication Co-authored-by: aider (openrouter/openai/o3-mini-high) --- amqp/rabbitmq/user_management_sample.py | 85 ++++++ .../user_management_service_client.py | 260 ++++++++++++++++ .../test_user_management_service_client.py | 284 ++++++++++++++++++ 3 files changed, 629 insertions(+) create mode 100644 amqp/rabbitmq/user_management_sample.py create mode 100644 amqp/rabbitmq/user_management_service_client.py create mode 100644 tests/rabbitmq/test_user_management_service_client.py diff --git a/amqp/rabbitmq/user_management_sample.py b/amqp/rabbitmq/user_management_sample.py new file mode 100644 index 0000000..b6a2c72 --- /dev/null +++ b/amqp/rabbitmq/user_management_sample.py @@ -0,0 +1,85 @@ +import asyncio +import logging +from typing import Dict, Any + +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()) diff --git a/amqp/rabbitmq/user_management_service_client.py b/amqp/rabbitmq/user_management_service_client.py new file mode 100644 index 0000000..a33c0de --- /dev/null +++ b/amqp/rabbitmq/user_management_service_client.py @@ -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} + ) diff --git a/tests/rabbitmq/test_user_management_service_client.py b/tests/rabbitmq/test_user_management_service_client.py new file mode 100644 index 0000000..21c6b59 --- /dev/null +++ b/tests/rabbitmq/test_user_management_service_client.py @@ -0,0 +1,284 @@ +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 + mock_channel = MagicMock() + 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()