diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..44dad06 --- /dev/null +++ b/.dockerignore @@ -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 diff --git a/.forgejo/workflows/build-and-publish.yml b/.forgejo/workflows/build-and-publish.yml new file mode 100644 index 0000000..abb7269 --- /dev/null +++ b/.forgejo/workflows/build-and-publish.yml @@ -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 diff --git a/.forgejo/workflows/publish-docker.yaml b/.forgejo/workflows/publish-docker.yaml index ea503a6..8461141 100644 --- a/.forgejo/workflows/publish-docker.yaml +++ b/.forgejo/workflows/publish-docker.yaml @@ -2,7 +2,7 @@ on: [push] env: REGISTRY_URL: "git.cleverthis.com" - REPOSITORY: "clevermicro/amq-adapter-python" + REPOSITORY: "clevermicro/amq-adapter-python-demo" DOCKER_HOST: "tcp://dind:2375" jobs: @@ -38,4 +38,4 @@ jobs: file: ./Dockerfile no-cache: true push: true - tags: ${{ env.REGISTRY_URL }}/${{ env.REPOSITORY }}:latest + tags: ${{ env.REGISTRY_URL }}/${{ env.REPOSITORY }}:20250709-01 diff --git a/.gitignore b/.gitignore index 9f169a6..abbebc3 100644 --- a/.gitignore +++ b/.gitignore @@ -56,7 +56,8 @@ amq_adapter.egg-info/ build/ unused-code/ - +.vscode/ amq_adapter.egg-info/ build/ coverage.xml +.aider* diff --git a/Dockerfile b/Dockerfile index f1efb0c..8af4c28 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,24 +1,32 @@ -# Use the official Python image as the base image -FROM python:3.11-slim - -# Set the working directory to /app +# Build stage +FROM python:3.11-slim AS builder WORKDIR /app ENV PYTHONPATH=/app -# Install system dependencies -RUN apt-get update && apt-get install -y \ - gcc \ - && rm -rf /var/lib/apt/lists/* +# Install system dependencies for building +RUN apt-get update && apt-get install -y gcc && rm -rf /var/lib/apt/lists/* -# Copy requirements and install Python dependencies +# Copy project files for building COPY pyproject.toml . -RUN pip install --no-cache-dir build -RUN pip install --no-cache-dir . +COPY . . -# Install additional dependencies for demo.py +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.9 \ + python-multipart==0.0.20 \ httpx==0.27.0 \ opentelemetry-api==1.34.1 \ opentelemetry-sdk==1.34.1 \ @@ -28,7 +36,7 @@ RUN pip install --no-cache-dir \ opentelemetry-instrumentation-httpx==0.55b1 \ protobuf==5.28.3 -# Copy the application code +# Copy the application code if needed COPY ./amqp ./amqp COPY ./otdemo ./otdemo COPY ./demo.py . @@ -36,8 +44,6 @@ COPY ./demo.py . # Create directory for uploaded files RUN mkdir -p /tmp/clevermicro_documents -# Expose the port EXPOSE 8000 8001 -# Run the FastAPI application CMD ["uvicorn", "demo:app", "--host", "0.0.0.0", "--port", "8000"] diff --git a/README.md b/README.md index 8e9b5ff..36ba7ce 100644 --- a/README.md +++ b/README.md @@ -116,3 +116,36 @@ AMQ_VERBOSE_LOGGING - if set to 1, the adapter will indicate in the log next to As an open-source project we run entierly off donations. Buy one of our hardworking developers a beer by [donating at OpenCollective](https://opencollective.com/cleverthis). All donations go to our bounty fund and allow us to place bounties on important bugs and enhancements. You may also use Beerpay linked via the above badges. +## Running with Docker Compose + +This project includes a Docker Compose configuration for easy local development and testing. + +### Prerequisites + +- Docker and Docker Compose installed on your system + +### Starting the services + +```bash +docker-compose up -d +``` + +This will start: +- The demo application (accessible at http://localhost:8000) +- Jaeger for distributed tracing (UI at http://localhost:16686) +- Prometheus for metrics collection (UI at http://localhost:9090) +- Alertmanager for alert handling (UI at http://localhost:9093) +- RabbitMQ for messaging (Management UI at http://localhost:15672) + +### Stopping the services + +```bash +docker-compose down +``` + +To remove volumes as well: + +```bash +docker-compose down -v +``` + diff --git a/amqp/adapter/backpressure_handler.py b/amqp/adapter/backpressure_handler.py index 97fd5a0..478387f 100644 --- a/amqp/adapter/backpressure_handler.py +++ b/amqp/adapter/backpressure_handler.py @@ -3,8 +3,8 @@ import json import time from asyncio import AbstractEventLoop from threading import Thread +from typing import Optional -import psutil from aio_pika import Message from aio_pika.abc import AbstractRobustChannel @@ -20,15 +20,15 @@ class ScaleRequestV1: self, serviceId: str, taskId: str, - maxAvailability: int, - currentAvailability: int, + max_availability: int, + current_availability: int, requestType: ScalingRequestAlert, ): self.version = 1 # Version of the request, currently 1 self.serviceId = serviceId self.taskId = taskId - self.maxAvailability = maxAvailability - self.currentAvailability = currentAvailability + self.max_availability = max_availability + self.current_availability = current_availability self.requestType = requestType self.thread = None @@ -39,60 +39,18 @@ class ScaleRequestV1: "version": self.version, "serviceId": self.serviceId, "taskId": self.taskId, - "maxAvailability": self.maxAvailability, - "currentAvailability": self.currentAvailability, + "maxAvailability": self.max_availability, + "currentAvailability": self.current_availability, "requestType": self.requestType.value, # Using .value for Enum serialization }, indent=2, ) -class RunningAverage: - """ - A class to maintain a running average of the last N values. - The intended use is to track CPU usage over time window, that is directly related to sample rate and window size. - This is the reason why the constructor takes the number of items to store, calculated as window size / sample rate - Initial_value is for unit test to emulate OVERLOAD or IDLE condition. - """ - - def __init__(self, num_items, initial_value=0): - self.buffer = [initial_value] * num_items # Fixed-size array - self.pointer = 0 # Current position in array - self.num_items = num_items # Maximum number of items to store - self.is_filled = False # Track if buffer is fully filled - - def add(self, value): - # Store the value at current pointer position - self.buffer[self.pointer] = value - - # Move pointer to next position with wrap-around - self.pointer += 1 - if self.pointer >= self.num_items: - self.pointer = 0 - self.is_filled = True - - def get_average(self): - # Determine how many items we should average - count = self.num_items if self.is_filled else self.pointer - - if count == 0: - return 0.0 # Avoid division by zero - - return sum(self.buffer) / count - - def get_current_values(self): - """Returns the values in chronological order (oldest first)""" - if not self.is_filled: - return self.buffer[: self.pointer] - - return self.buffer[self.pointer :] + self.buffer[: self.pointer] - - class BackpressureHandler: # Track the number of messages currently being processed - current_parallel_executions = 0 + current_availability = 0 # helps detect IDLE condition - last_data_message_time = 0 # helps prevent flooding the system with backpressure events last_backpressure_event_time = 0 last_backpressure_event = ScalingRequestAlert.UPDATE @@ -112,48 +70,64 @@ class BackpressureHandler: self.exchange = None self.swarm_service_id = self.config.amq_adapter.swarm_service_id self.swarm_task_id = self.config.amq_adapter.swarm_task_id - self.parallel_workers = self.config.backpressure.threshold_threads - self.average_cpu_usage = None self.do_loop = -1 + self._resource_usage_changed = 0 + self._resource_average_value = 0 + self._last_resource_max_value = 100 # Default max value for CPU usage + self.max_availability = -1 + self.current_availability = 1 + self.max_load = 1 + self.current_load = 0 - def _do_loop(self) -> bool: - """ - Check if the loop should continue running. - Positive value of do_loop indicates the number of iterations left. - Negative value indicates the loop should run indefinitely. - """ - _val = self.do_loop != 0 - if self.do_loop > 0: - self.do_loop -= 1 - return _val - - def increase_parallel_executions(self): - """Increase the number of parallel executions""" + def increase_current_load(self): + """Increase the number of parallel executions, or current_availability load""" logging_info( - "Backpressure: Increase parallel executions, current=%s", - self.current_parallel_executions, + "Backpressure: Increase current load (%s) by 1", + self.current_load, ) - self.current_parallel_executions += 1 + self.current_load += 1 - def decrease_parallel_executions(self): - """Decrease the number of parallel executions""" + def decrease_current_load(self): + """Decrease the number of parallel executions, or current load""" logging_info( - "Backpressure: Decrease parallel executions, current=%s", - self.current_parallel_executions, + "Backpressure: Decrease current load(%s) by 1", + self.current_load, ) - if self.current_parallel_executions > 0: - self.current_parallel_executions -= 1 + if self.current_load > 0: + self.current_load -= 1 - def update_parallel_executions(self, count: int): - """Update the number of parallel executions""" - self.current_parallel_executions = count + def update_current_availability(self, count: int): + """Update the number of parallel executions, or current load""" + self.current_availability = count - def update_last_data_message_time(self): - logging_info("Backpressure: Update last data message time") + def update_last_backpressure_event_time(self): + # logging_info("Backpressure: Update last data message time") """Update the last data message time""" - self.last_data_message_time = time.time() + self.last_backpressure_event_time = time.time() - def start_backpressure_monitor(self) -> Thread: + async def update_backpressure_value(self, current_availability: int, maximum: int): + """ + Update the current backpressure value and check for overload conditions. + This method is called by the AMQService.backpressure() method to update + the current parallel executions count and check for overload conditions. + + Args: + current_availability: Current value of the backpressure metric + maximum: Maximum value of the backpressure metric + """ + logging_info( + "Backpressure: Updating backpressure value, current_availability=%s, maximum=%s", + current_availability, + maximum, + ) + if maximum > 0 and maximum > self.max_availability: + self.max_availability = maximum + # Update the current_availability / parallel executions count + self.update_current_availability(current_availability) + # Check for overload conditions + await self.check_overload_condition() + + def start_backpressure_monitor(self) -> Optional[Thread]: # Start the Backpressure monitor loop self.thread = Thread(target=self.backpressure_monitor_loop) self.thread.daemon = True # This makes it a daemon thread @@ -161,103 +135,90 @@ class BackpressureHandler: return self.thread async def check_overload_condition(self): - """Check if the current parallel executions exceed the limit""" - self.update_last_data_message_time() + """ + Check if the current_availability availability is too low (OVERLOAD) + or if the service has been idle for too long with high availability (IDLE). + + Note: current_availability represents available capacity, not used capacity. + - Low availability (close to 0) means OVERLOAD + - High availability (close to maximum) with no activity means IDLE + """ + current_time = time.time() + last_event_delta: float = current_time - self.last_backpressure_event_time + logging_info( - "Backpressure: Check overload condition, current=%s, max=%s", - self.current_parallel_executions, - self.parallel_workers, + "Backpressure: Check conditions, current_availability=%s, max=%s, last_activity=%s", + self.current_availability, + self.max_availability, + last_event_delta, ) - if self.current_parallel_executions >= self.parallel_workers - 1: + + # Check for OVERLOAD condition - low availability (less than 20% of maximum) + if self.current_availability <= round(0.1 * self.max_availability): # Check if the last backpressure event was not an overload - if self.last_backpressure_event != ScalingRequestAlert.OVERLOAD: - # Trigger the overload event - await self.handle_backpressure_overload_event() - self.last_backpressure_event_time = time.time() - self.last_backpressure_event = ScalingRequestAlert.OVERLOAD - - async def _backpressure_monitor(self): - """Monitor the backpressure conditions and trigger events accordingly""" - _time_window_millis = self.config.backpressure.time_window - _idle_duration_millis = self.config.backpressure.idle_duration - _overload_duration_millis = self.config.backpressure.cpu_overload_duration - _monitor_interval = 0.5 # Monitor every 500ms - CPU_IDLE, CPU_ACTIVE, CPU_OVERLOAD = 0, 1, 2 - IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 15, 90 - if self.average_cpu_usage is None: - self.average_cpu_usage = RunningAverage( - int(max(_overload_duration_millis, _idle_duration_millis) // _monitor_interval) - ) - _cpu_usage_state = CPU_ACTIVE - _cpu_usage_changed = time.time() - while self._do_loop(): - _current_cpu_usage = psutil.cpu_percent(interval=None) - self.average_cpu_usage.add(_current_cpu_usage) - _average_cpu_usage = self.average_cpu_usage.get_average() - _now = time.time() - # logging_info( - # "Backpressure: current CPU[pct]=%s, avg=%s, state=%s, last dataMsg=%s, last event=%s, last eventTime=%s", - # _current_cpu_usage, - # _average_cpu_usage, - # _cpu_usage_state, - # self.last_data_message_time, - # self.last_backpressure_event, - # self.last_backpressure_event_time, - # ) - if _average_cpu_usage < IDLE_THRESHOLD: - if _cpu_usage_state != CPU_IDLE: - _cpu_usage_changed = _now - _cpu_usage_state = CPU_IDLE - elif _average_cpu_usage > OVERLOAD_THRESHOLD: - if _cpu_usage_state != CPU_OVERLOAD: - _cpu_usage_changed = _now - _cpu_usage_state = CPU_OVERLOAD - else: - if _cpu_usage_state != CPU_ACTIVE: - _cpu_usage_changed = _now - _cpu_usage_state = CPU_ACTIVE - - # Check if the current time exceeds the last backpressure event time if ( - _cpu_usage_state == CPU_OVERLOAD - and _now - _cpu_usage_changed > _overload_duration_millis - and ( - _now - self.last_backpressure_event_time > _time_window_millis - or self.last_backpressure_event != ScalingRequestAlert.OVERLOAD - ) + self.last_backpressure_event != ScalingRequestAlert.OVERLOAD + or last_event_delta > self.config.backpressure.idle_duration ): # Trigger the overload event await self.handle_backpressure_overload_event() - self.last_backpressure_event_time = _now + # update / reset time-window so that the OVERLOAD is not sent too often + self.last_backpressure_event_time = current_time self.last_backpressure_event = ScalingRequestAlert.OVERLOAD - elif ( - _cpu_usage_state == CPU_IDLE - and _now - _cpu_usage_changed > _idle_duration_millis - and _now - self.last_data_message_time > _idle_duration_millis - and ( - _now - self.last_backpressure_event_time > _time_window_millis - or self.last_backpressure_event != ScalingRequestAlert.IDLE - ) + + # Check for IDLE condition - high availability (more than 80% of maximum) with no activity + elif self.current_availability >= round(0.8 * self.max_availability): + idle_duration = self.config.backpressure.idle_duration + # Check if service has been idle for longer than the configured duration + if ( + last_event_delta > idle_duration + and self.last_backpressure_event == ScalingRequestAlert.IDLE ): + logging_info( + "Backpressure: Service has been idle for %s seconds (threshold: %s seconds)", + last_event_delta, + idle_duration, + ) + # Trigger the idle event await self._handle_backpressure_idle_event() - self.last_backpressure_event = ScalingRequestAlert.IDLE - self.last_backpressure_event_time = _now + # update / reset time-window so that the IDLE is not sent too often + self.last_backpressure_event_time = current_time else: - # Trigger update event - if _now - self.last_backpressure_event_time > _time_window_millis: - # Trigger the update event - _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 + # Don't send IDLE right away when the availability increases, but wait for a while, send UPDATE instead + self.last_backpressure_event = ScalingRequestAlert.IDLE + _scaling_request: ScaleRequestV1 = ScaleRequestV1( + self.swarm_service_id, + self.swarm_task_id, + 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 + + self.update_last_backpressure_event_time() + + async def _backpressure_monitor(self): + """Periodically monitor the backpressure conditions and trigger events accordingly""" + _monitor_interval = 0.5 # Monitor every 500ms + while self._do_loop(): + await self.check_overload_condition() await asyncio.sleep(_monitor_interval) def backpressure_monitor_loop(self): @@ -270,14 +231,12 @@ class BackpressureHandler: scaling_request: ScaleRequestV1 = ScaleRequestV1( self.swarm_service_id, self.swarm_task_id, - self.parallel_workers, - self.parallel_workers - self.current_parallel_executions, + self.max_availability, + self.current_availability, # Current availability is passed directly ScalingRequestAlert.OVERLOAD, ) # Address the message to any management service instance capable of supporting BACKPRESSURE request await self.publish_backpressure_request(scaling_request) - # update / reset time-window so that the OVERLOAD is not sent too often - self.last_data_message_time = time.time() async def _handle_backpressure_idle_event(self): logging_warning("Backpressure: Service is idle.") @@ -285,14 +244,12 @@ class BackpressureHandler: scaling_request: ScaleRequestV1 = ScaleRequestV1( self.swarm_service_id, self.swarm_task_id, - self.parallel_workers, - self.parallel_workers - self.current_parallel_executions, + self.max_availability, + self.current_availability, # Current availability is passed directly ScalingRequestAlert.IDLE, ) # Address the message to any adapter capable of supporting BACKPRESSURE request await self.publish_backpressure_request(scaling_request) - # update / reset time-window so that the IDLE is not sent too often - self.last_data_message_time = time.time() async def publish_backpressure_request(self, scaling_request: ScaleRequestV1): # Publish the backpressure request to the management service @@ -339,3 +296,14 @@ class BackpressureHandler: if BackpressureHandler._callback_list is not None: BackpressureHandler._callback_list["_wrap_rabbit_mq_api"] = 1 await await_future(asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api(), self.loop)) + + def _do_loop(self) -> bool: + """ + Helper function for unit tests to perform several loops only. Check if the loop should continue running. + Positive value of do_loop indicates the number of iterations left. + Negative value indicates the loop should run indefinitely. + """ + _val = self.do_loop != 0 + if self.do_loop > 0: + self.do_loop -= 1 + return _val diff --git a/amqp/adapter/cleverthis_service_adapter.py b/amqp/adapter/cleverthis_service_adapter.py index 3e4c66b..2a145b5 100644 --- a/amqp/adapter/cleverthis_service_adapter.py +++ b/amqp/adapter/cleverthis_service_adapter.py @@ -83,7 +83,7 @@ async def _convert_file_response( temporary dedicated queue. """ # 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 + # to parent event loop, and need to execute at that event loop, not the current_availability one # allow coroutine to execute. calling _future.result() will block the entire thread # and the coroutine will never execute _channel, _exchange, _routing_key = await await_result( @@ -241,7 +241,7 @@ class CleverThisServiceAdapter: async def get_current_user(self, token: str) -> object: """ - To Be Overriden in subclass for given CleverThis Service, to return the current user based + To Be Overriden in subclass for given CleverThis Service, to return the current_availability user based on the token, in format appropriate for the service. :param token: The token to be used for authentication. @@ -491,6 +491,9 @@ class CleverThisServiceAdapter: _routes = self.routes.get(method, []) _path, _args = AMQDataParser.parse_get_url(message.path()) for _route in _routes: + # TODO: here the regex will match with path prefix, + # but it might not be a stable feature that we could rely on. + # Need to figure out something stable for the prefix path. _match = _route.path_regex.match(_path) if _match: _args.update(_match.groupdict()) diff --git a/amqp/adapter/consul_global_id_generator.py b/amqp/adapter/consul_global_id_generator.py new file mode 100644 index 0000000..9a5e893 --- /dev/null +++ b/amqp/adapter/consul_global_id_generator.py @@ -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() diff --git a/amqp/adapter/data_message_factory.py b/amqp/adapter/data_message_factory.py index b42f011..02a17ad 100644 --- a/amqp/adapter/data_message_factory.py +++ b/amqp/adapter/data_message_factory.py @@ -7,10 +7,11 @@ import json import re from datetime import datetime, timezone from io import BytesIO -from typing import Dict, List +from typing import Any, Dict, List import amqp.adapter.file_handler from amqp.adapter.file_handler import FileHandler, StreamingFileHandler +from amqp.adapter.pydantic_serializer import python_type_to_json from amqp.adapter.serializer import ( map_as_string, map_with_list_as_string, @@ -112,6 +113,32 @@ class DataMessageFactory: payload, ) + def create_rpc_message( + self, + fq_method_name: str, + swarm_id: str, + args: Dict[str, Any] = None, + ) -> DataMessage: + """ + Create a RPC message with the given parameters. + :param fq_method_name: fully qualified method name (e.g. 'com.example.ServiceClass.method_name') + :param args: arguments whose serialized form will form the request body + :return: DataMessage object + """ + package, class_name, method = fq_method_name.rsplit(".", 2) + headers = {"Content-Type": ["application/octet-stream"], "Return-Address": [swarm_id]} + trace_info = {} + return DataMessage( + DataMessageMagicByte.RPC_REQUEST.value, + self.generate_snowflake_id(), + package, + class_name, + method, + headers, + trace_info, + python_type_to_json(args) if args else b"", + ) + @staticmethod def amq_message_routing_key(message: DataMessage) -> str: """ diff --git a/amqp/adapter/logging_utils.py b/amqp/adapter/logging_utils.py index d8ff2d8..0238f8f 100644 --- a/amqp/adapter/logging_utils.py +++ b/amqp/adapter/logging_utils.py @@ -13,17 +13,17 @@ __verbose__ = os.environ.get("AMQ_VERBOSE_LOGGING", "0") == "1" def get_context_info(): """ - Retrieves the current module, line number, and function name. + Retrieves the current_availability module, line number, and function name. Returns: A tuple containing (module name, line number, function name). Returns (None, None, None) if it fails to retrieve the information. """ try: - # Get the frame object for the current function call + # Get the frame object for the current_availability function call frame = inspect.currentframe() if frame is not None: - # Get the frame object for the caller of the current function + # Get the frame object for the caller of the current_availability function frame = frame.f_back.f_back if frame: code_context = frame.f_code @@ -44,7 +44,7 @@ def get_context_info(): def logging_error(msg: str, *args) -> None: """ - Log an error message according to current logging for 'ERROR' level. + Log an error message according to current_availability logging for 'ERROR' level. Add module name, line and function name where the error occurred, on single line. Args: msg (str): The error message to log. @@ -70,7 +70,7 @@ def logging_error(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. Args: 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: """ - Log an info message according to current logging for 'INFO' level. + Log an info message according to current_availability logging for 'INFO' level. Add module name, line and function name where the print occurred, on single line. Args: msg (str): The info message to log. @@ -112,7 +112,7 @@ def logging_info(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. Args: msg (str): The debug message to log. diff --git a/amqp/config/amq_configuration.py b/amqp/config/amq_configuration.py index 383c7ef..79dd4e3 100644 --- a/amqp/config/amq_configuration.py +++ b/amqp/config/amq_configuration.py @@ -24,9 +24,7 @@ class AMQAdapter: :param config: config parser to read configuration values """ - self.generator_id = config.getint( - "CleverMicro-AMQ", self.PREFIX + "generator-id", fallback=164 - ) + self.generator_id = 0 self.service_name = config.get( "CleverMicro-AMQ", self.PREFIX + "service-name", @@ -40,9 +38,72 @@ class AMQAdapter: self.PREFIX + "require-authenticated-user", fallback=False, ) + self.backpressure_enabled = config.getboolean( + "CleverMicro-AMQ", + self.PREFIX + "default-backpressure-enabled", + fallback=False, + ) # these 2 shall be provided by Docker swarm, format is a random alphanumeric string self.swarm_service_id = os.getenv("SWARM_SERVICE_ID", default="clever-amqp-service") self.swarm_task_id = os.getenv("SWARM_TASK_ID", default="python-adapter") + # + # Consul related configuration + self.consul_port = ( + int( + os.getenv( + "CONSUL_PORT", + config.get("CleverMicro-AMQ", self.PREFIX + "consul-port", fallback="8500"), + ) + ) + if os.getenv("CONSUL_PORT") is not None + else config.getint("CleverMicro-AMQ", self.PREFIX + "consul-port", fallback=8500) + ) + self.consul_host = os.getenv( + "CONSUL_HOST", + config.get("CleverMicro-AMQ", self.PREFIX + "consul-host", fallback="consul"), + ) + self.consul_max_retries = ( + int( + os.getenv( + "CONSUL_MAX_RETRIES", + config.get("CleverMicro-AMQ", self.PREFIX + "consul-max-retries", fallback="5"), + ) + ) + if os.getenv("CONSUL_MAX_RETRIES") is not None + else config.getint("CleverMicro-AMQ", self.PREFIX + "consul-max-retries", fallback=5) + ) + self.consul_initial_retry_delay = ( + float( + os.getenv( + "CONSUL_INITIAL_RETRY_DELAY", + config.get( + "CleverMicro-AMQ", + self.PREFIX + "consul-initial-retry-delay", + fallback="0.2", + ), + ) + ) + if os.getenv("CONSUL_INITIAL_RETRY_DELAY") is not None + else config.getfloat( + "CleverMicro-AMQ", self.PREFIX + "consul-initial-retry-delay", fallback=0.2 + ) + ) + self.consul_counter_key = ( + os.getenv( + "CONSUL_COUNTER_KEY", + config.get( + "CleverMicro-AMQ", + self.PREFIX + "consul-counter-key", + fallback="service/ids/counter", + ), + ) + if os.getenv("CONSUL_COUNTER_KEY") is not None + else config.get( + "CleverMicro-AMQ", + self.PREFIX + "consul-counter-key", + fallback="service/ids/counter", + ) + ) def adapter_prefix(self) -> str: return f"{self.service_name}-{self.generator_id}-" @@ -102,7 +163,7 @@ class Backpressure: :param config: config parser to read configuration values """ - self.threshold_threads = config.getint( + self.threshold_load = config.getint( "CleverMicro-AMQ", self.PREFIX + "threshold", fallback=int(os.getenv("PARALLEL_WORKERS", 1)), @@ -124,6 +185,10 @@ class Backpressure: self.idle_duration = config.getint( "CleverMicro-AMQ", self.PREFIX + "idle-duration", fallback=30 ) + # Flag to enable/disable CPU monitoring + self.cpu_monitoring_enabled = config.getboolean( + "CleverMicro-AMQ", self.PREFIX + "cpu-monitoring-enabled", fallback=False + ) class AMQConfiguration: @@ -138,7 +203,7 @@ class AMQConfiguration: config.read(config_file_name) else: # Get the directory of the caller (the module where this function is called) - caller_frame = inspect.stack()[1] # 0 is current frame, 1 is caller + caller_frame = inspect.stack()[1] # 0 is current_availability frame, 1 is caller caller_module = inspect.getmodule(caller_frame[0]) if caller_module: caller_dir = os.path.dirname(os.path.abspath(caller_module.__file__)) @@ -163,3 +228,14 @@ class AMQConfiguration: self.amq_adapter: AMQAdapter = AMQAdapter(config) self.dispatch: Dispatch = Dispatch(config) self.backpressure: Backpressure = Backpressure(config) + + def get_unique_instance_id(self) -> int: + """ + Get a unique instance ID for this AMQ service. + + Returns: + int: A unique instance ID + """ + if self.instance_id == -1: + logging_warning(f"Generated unique instance ID: {self.instance_id}") + return self.instance_id diff --git a/amqp/config/application.properties b/amqp/config/application.properties index fdfa5a8..7a0c3c7 100644 --- a/amqp/config/application.properties +++ b/amqp/config/application.properties @@ -14,6 +14,12 @@ cm.amq-adapter.generator-id=1 cm.amq-adapter.route-mapping= # Indicate if AMQ Adapter should respond with HTTP 401 Unauthorized if the request does not contain valid JWT token cm.amq-adapter.require-authenticated-user=false +# Consul configuration for CleverMicro AMQ Adapter +cm.amq-adapter.consul-host=consul +cm.amq-adapter.consul-port=8500 +cm.amq-adapter.consul-max-retries=5 +cm.amq-adapter.consul-initial-retry-delay=0.2 +cm.amq-adapter.consul-id-generator-key=services/ids/counter # CleverMicro AMQ Adapter dispatcher, settings related to RabbitMQ connection and optional REST forwarding cm.dispatch.use-dlq=true @@ -27,7 +33,8 @@ cm.dispatch.download-dir=/tmp/downloaded_files #cm.dispatch.service-port=8080 # Backpressure monitor settings -# Define maximum number of messages of being processed in parallel before triggering backpressure OVERLOAD alert +# Define maximum number of resources that service can use before triggering backpressure OVERLOAD alert +# 'resource' is a generic term that can mean CPU, memory, parallel requests, etc. cm.backpressure.threshold=5 # Define maximum CPU usage (%) before triggering backpressure OVERLOAD alert cm.backpressure.threshold-cpu-overload=90 @@ -39,4 +46,5 @@ cm.backpressure.time-window=10 cm.backpressure.cpu-overload-duration=30 # Define the time window length, in seconds, during which, if no request is made to the Service, the system will trigger IDLE alert cm.backpressure.cpu-idle-duration=30 - +# AMQ Adapter can monitor CPU usage and trigger backpressure alerts based on CPU usage. Enable this feature explicitly +cm.backpressure.cpu-monitoring-enabled=false diff --git a/amqp/model/model.py b/amqp/model/model.py index fe31fc4..08742b3 100644 --- a/amqp/model/model.py +++ b/amqp/model/model.py @@ -91,6 +91,9 @@ class CleverMicroMessage: def content_type(self) -> str: return self.headers().get("Content-Type", self.DEFAULT_CONTENT_TYPE)[0] + def return_address(self) -> str: + return self.headers().get("Return-Address", "")[0] + class DataMessage(CleverMicroMessage): """ diff --git a/amqp/rabbitmq/rabbit_mq_client.py b/amqp/rabbitmq/rabbit_mq_client.py index 3804a84..d839ce6 100644 --- a/amqp/rabbitmq/rabbit_mq_client.py +++ b/amqp/rabbitmq/rabbit_mq_client.py @@ -17,7 +17,6 @@ from amqp.router.router_base import ( AMQ_RPC_EXCHANGE, DLQ_EXCHANGE, ) -from amqp.router.router_consumer import RouterConsumer from amqp.router.utils import sanitize_as_rabbitmq_name @@ -25,7 +24,7 @@ class RabbitMQClient: PREFETCH_COUNT = 1 DLQ_ARGS = {"x-dead-letter-exchange": DLQ_EXCHANGE} - def __init__(self, configuration): + def __init__(self, configuration, router): self.channel: aio_pika.RobustChannel | None = None self.connection: aio_pika.RobustConnection | None = None self.rpc_exchange: aio_pika.RobustExchange | None = None @@ -37,7 +36,7 @@ class RabbitMQClient: self.amq_configuration.dispatch.amq_host, ) logging.info("*******************************************") - self.router = RouterConsumer(configuration) + self.router = router # this.meterRegistry = new PrometheusMeterRegistry(PrometheusConfig.DEFAULT); self.reply_to_exchange: AbstractRobustExchange | None = ( 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/amqp/router/route_database.py b/amqp/router/route_database.py index 4f36d68..56406c6 100644 --- a/amqp/router/route_database.py +++ b/amqp/router/route_database.py @@ -1,5 +1,5 @@ import re -from typing import Set +from typing import Optional, Set from amqp.adapter.logging_utils import logging_info from amqp.model.model import AMQRoute @@ -55,6 +55,13 @@ class RouteDatabase: return route return None + def find_route_by_service(self, service_name: str) -> Optional[AMQRoute]: + logging_info(f"findRoute service_name:{service_name} in:{self.wrap_routes()}") + for route in self.routes: + if route.component_name == service_name: + return route + return None + def add_route(self, route: AMQRoute) -> bool: if route not in self.routes: self.routes.add(route) diff --git a/amqp/router/router_base.py b/amqp/router/router_base.py index dd760f1..f7cc4ab 100644 --- a/amqp/router/router_base.py +++ b/amqp/router/router_base.py @@ -1,4 +1,4 @@ -from typing import Callable, List, Set +from typing import Callable, List, Optional, Set from amqp.adapter.amq_route_factory import NULL_ROUTE, AMQRouteFactory from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info @@ -49,6 +49,9 @@ class RouterBase: def find_route_by_message(self, message: DataMessage) -> AMQRoute | None: return self.route_database.find_route(message.domain(), message.path()) + def find_route_by_service(self, service_name: str) -> Optional[AMQRoute]: + return self.route_database.find_route_by_service(service_name) + def add_routes(self, message: str) -> None: try: amq_route_list = self.unwrap_route_list(message) diff --git a/amqp/router/router_consumer.py b/amqp/router/router_consumer.py index 10f0a23..afea07e 100644 --- a/amqp/router/router_consumer.py +++ b/amqp/router/router_consumer.py @@ -174,12 +174,12 @@ class RouterConsumer(RouterProducer): ) if not await self.publish_service_message(service_message, self.cm_response_exchange): logging_warning( - "RouterConsumer couldn't remove current route: %s, channel not open.", + "RouterConsumer couldn't remove current_availability route: %s, channel not open.", route, ) except Exception as ioe: - logging_error("RouterConsumer failed remove current route: %s", ioe) + logging_error("RouterConsumer failed remove current_availability route: %s", ioe) result: bool = self.remove_route(route) return result diff --git a/amqp/router/router_producer.py b/amqp/router/router_producer.py index 6224dc1..2dd8bd6 100644 --- a/amqp/router/router_producer.py +++ b/amqp/router/router_producer.py @@ -81,7 +81,7 @@ class RouterProducer(RouterBase): name=CM_REQUEST_EXCHANGE, type=ExchangeType.fanout ) # - # 2. declare ResponseExchange where to publish current RoutingData (response to + # 2. declare ResponseExchange where to publish current_availability RoutingData (response to # a RoutingData request) self.cm_response_exchange = await self.channel.declare_exchange( name=CM_RESPONSE_EXCHANGE, type=ExchangeType.fanout diff --git a/amqp/service/amq_message_handler.py b/amqp/service/amq_message_handler.py index af67d0d..8dea0e4 100644 --- a/amqp/service/amq_message_handler.py +++ b/amqp/service/amq_message_handler.py @@ -123,16 +123,16 @@ class DataMessageHandler: :param loop: Event loop :return: DataResponse """ - # Increment the counter for parallel executions - self.backpressure_handler.increase_parallel_executions() + # Increment the counter for parallel executions and check resource availability + self.backpressure_handler.increase_current_load() + await self.backpressure_handler.check_overload_condition() logging_info( - f"PARALLEL: Current Capacity [{self.backpressure_handler.current_parallel_executions}] of [{self.backpressure_handler.parallel_workers}]" + f"PARALLEL: Current Capacity [{self.backpressure_handler.current_availability}] of [{self.backpressure_handler.max_availability}]" ) if message.reply_to: self.reply_to = message.reply_to try: - await self.backpressure_handler.check_overload_condition() _a_message = AMQMessage(amq_message, self.rabbit_mq_client.connection) _response: DataResponse = await self.service_adapter.on_message(_a_message) logging_debug(f"Result in thread: {_response}") @@ -155,7 +155,7 @@ class DataMessageHandler: except Exception as e: logging_error(f"Request handler: {e}") - self.backpressure_handler.decrease_parallel_executions() + self.backpressure_handler.decrease_current_load() async def reconstitute_data_message( self, message: AbstractIncomingMessage @@ -207,12 +207,12 @@ class DataMessageHandler: # map_as_string(amq_message.trace_info()), context_with_span, context) def record_duration(self, start, delivery_tag): - global last_data_message_time - last_data_message_time = time.time() + global last_backpressure_event_time + last_backpressure_event_time = time.time() logging_info( "######[%s] Done, single ACK: %sms.", delivery_tag, - last_data_message_time - start, + last_backpressure_event_time - start, ) async def send_reply(self, reply_to: str, response: DataResponse): diff --git a/amqp/service/amq_service.py b/amqp/service/amq_service.py index 0e161ee..d7b82d0 100644 --- a/amqp/service/amq_service.py +++ b/amqp/service/amq_service.py @@ -10,12 +10,14 @@ from aio_pika.abc import AbstractRobustExchange from amqp.adapter.amq_route_factory import AMQRouteFactory from amqp.adapter.backpressure_handler import BackpressureHandler +from amqp.adapter.consul_global_id_generator import get_unique_instance_id from amqp.adapter.data_message_factory import DataMessageFactory from amqp.adapter.logging_utils import logging_info, logging_warning from amqp.adapter.service_message_factory import ServiceMessageFactory from amqp.config.amq_configuration import AMQConfiguration -from amqp.model.model import DataMessage +from amqp.model.model import AMQRoute, DataMessage from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient +from amqp.router.router_consumer import RouterConsumer from amqp.service.amq_message_handler import DataMessageHandler from amqp.service.tracing import TracingConfig, initialize_telemetry @@ -36,16 +38,20 @@ class AMQService: def __init__(self, amq_configuration: AMQConfiguration, service_adapter=None): logging.info("***********************************************") - logging.info("********** AMQ Service *******************") + logging.info("************** AMQ Service ***************") logging.info("***********************************************") self.loop = asyncio.new_event_loop() self.amq_configuration = amq_configuration self.service_adapter = service_adapter + self.unique_id = get_unique_instance_id(amq_configuration.amq_adapter) + amq_configuration.amq_adapter.generator_id = self.unique_id + self.data_message_factory = DataMessageFactory.get_instance( - self.amq_configuration.amq_adapter.generator_id + self.unique_id & DataMessageFactory.MACHINE_ID_MASK ) - self.rabbit_mq_client = RabbitMQClient(self.amq_configuration) + 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, @@ -57,8 +63,18 @@ class AMQService: self.amq_configuration.amq_adapter.service_name ) self.backpressure_thread: Optional[Thread] = None + self.backpressure_handler: Optional[BackpressureHandler] = None + self.maximum_availability = -1 + # ======================================================================= + # ======================================================================= + # ============== P u b l i c A P I M e t h o d s =================== + # ======================================================================= + # ======================================================================= def run(self): + """ + Start the AMQ service, initializing RabbitMQ connection and message handlers. + """ # self.init(loop) will initialize RabbitMQ connection asyncio.set_event_loop(self.loop) self.loop.run_until_complete(self.init(self.loop, self.service_adapter)) @@ -68,9 +84,76 @@ class AMQService: logging.info("***********************************************") self.loop.run_forever() + def backpressure(self, current_availability: int, maximum: int = -1): + """ + Send Backpressure event to the Management Service. The Backpressure logic determines the type of the event + based on the current_availability and maximum values. + - OVERLOAD event is triggered immediately when the current_availability value is below threshold + typically set to 10% of the maximum value, or 0 if no maximum is provided. + - IDLE event is triggered when the current_availability value exeeds 80% of the maximum value or is greater + than 0 if no maximum is provided. + - UPDATE event is sent if neither OVERLOAD nor IDLE conditions are met, indicating that the backpressure is + within acceptable limits. + :param current_availability: Currently available capacity. + :param maximum: Maximum value of the backpressure metric. If undefined or set to lt 0 (e.g. -1), + the maximum is than determined from max value of current_availability seen over the time. + """ + if self.backpressure_handler is not None: + # watch maximum & current value to always have some valid maximum reference + if maximum < 0: + if current_availability > self.maximum_availability: + self.maximum_availability = current_availability + else: + if maximum > self.maximum_availability: + self.maximum_availability = maximum + asyncio.run_coroutine_threadsafe( + self.backpressure_handler.update_backpressure_value( + current_availability, maximum if maximum > 0 else self.maximum_availability + ), + loop=self.backpressure_handler.loop, + ) + + def rpc(self, service_name: str, fq_method_name: str, **kwargs) -> Future: + """ + Send a Remote Procedure Call (RPC) message to the AMQ service. + :param service_name: The service name to receive the RPC message. + :param fq_method_name: The fully qualified method name to execute the RPC message. + :param kwargs: Additional keyword arguments to be passed to the RPC method. + :return: A Future that resolves when the RPC response is received. + """ + route: Optional[AMQRoute] = self.router.find_route_by_service(service_name) + if route is not None: + message: DataMessage = self.data_message_factory.create_rpc_message( + fq_method_name=fq_method_name, + swarm_id=self.amq_configuration.amq_adapter.swarm_task_id, + **kwargs, + ) + result: Future = self.amq_data_message_handler.send_rpc(route, message) + self.set_outstanding_future(message, result) + return result + + result = asyncio.Future() + result.set_exception( + NameError(f"No route exists. Service '{service_name}' not found in routing table.") + ) + return result + + def get_unique_instance_id(self) -> int: + """ + Provides cluster-wide unique instance ID for this AMQ service instance. + This ID is used to identify the service instance in the cluster and is unique across all instances. + IDs are managed by shared Consul KV store, ensuring that even if multiple instances of the service are running, + they will each have unique ID. + """ + return self.unique_id + + # ======================================================================= + # ====================== Internal Methods ============================ + # ======================================================================= + async def init(self, loop, service_adapter): _reply_to_exchange: AbstractRobustExchange = await self.rabbit_mq_client.init(loop) - backpressure_handler = BackpressureHandler( + self.backpressure_handler = BackpressureHandler( channel=self.rabbit_mq_client.channel, loop=loop, config=self.amq_configuration, @@ -84,11 +167,12 @@ class AMQService: self.rabbit_mq_client, self.amq_configuration, loop=loop, - backpressure_handler=backpressure_handler, + backpressure_handler=self.backpressure_handler, ) await self.register_routes() await self.rabbit_mq_client.request_routes() - self.backpressure_thread = backpressure_handler.start_backpressure_monitor() + logging_info("Starting backpressure monitor thread") + self.backpressure_thread = self.backpressure_handler.start_backpressure_monitor() def cleanup(self): logging_info("AMQService.cleanup() rmqc=%s", self.rabbit_mq_client) @@ -97,7 +181,7 @@ class AMQService: async def reinitialize_if_disconnected(self): if not self.rabbit_mq_client.channel or self.rabbit_mq_client.channel.is_closed: - self.rabbit_mq_client = RabbitMQClient(self.amq_configuration) + self.rabbit_mq_client = RabbitMQClient(self.amq_configuration, self.router) await self.rabbit_mq_client.init(loop=self.loop) await self.register_routes() await self.rabbit_mq_client.request_routes() @@ -122,7 +206,7 @@ class AMQService: ) return None - def send_message(self, message: DataMessage, future: Future): + def set_outstanding_future(self, message: DataMessage, future: Future): self.amq_data_message_handler.outstanding[str(message.id())] = future def remove_outstanding_future(self, message_id: str): diff --git a/demo.py b/demo.py index 18b38eb..9d443b6 100644 --- a/demo.py +++ b/demo.py @@ -94,6 +94,7 @@ internal_api_call_duration_histogram = meter.create_histogram( 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 @@ -105,8 +106,6 @@ HTTPXClientInstrumentor().instrument() UPLOAD_DIR = "/tmp/clevermicro_documents" os.makedirs(UPLOAD_DIR, exist_ok=True) -otadapter = OTDemoAmqpAdapter(AMQConfiguration("otdemo/otdemo.properties"), app.routes) - @app.post("/api/v3/documents") async def upload_document( @@ -135,13 +134,13 @@ async def upload_document( start_time = time.time() async with httpx.AsyncClient() as client: - # httpx instrumentation ensures the current trace context is propagated + # 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://host.docker.internal: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) @@ -223,6 +222,8 @@ async def get_documents(): 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. @@ -232,7 +233,7 @@ async def get_documents(): # ``` # fastapi==0.111.0 # uvicorn[standard]==0.29.0 -# python-multipart==0.0.9 +# python-multipart==0.0.20 # httpx==0.27.0 # opentelemetry-api==1.25.0 # opentelemetry-sdk==1.25.0 diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..5315b02 --- /dev/null +++ b/docker-compose.yml @@ -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: diff --git a/otdemo/otdemo.properties b/otdemo/otdemo.properties index 5fb0120..e330510 100644 --- a/otdemo/otdemo.properties +++ b/otdemo/otdemo.properties @@ -5,13 +5,16 @@ # CleverMicro AMQ Adapter settings that identify this AMQ Adapter instance to other adapters # A name to identify this service for the AMQ Adapter -cm.amq-adapter.service-name=otdemo +cm.amq-adapter.service-name=cleverswarm # The CleverMicro AMQ Adapter instance ID, used to identify this instance in the CleverMicro AMQ Adapter cluster cm.amq-adapter.generator-id=1 # Route mapping. For route detail description see section 3.1 in Onboarding Guide for Python: # https://docs.cleverthis.com/en/architecture/microservices/onboarding/python # NOTE: this is a comma-separated list of route mappings (multiple, comma-separated routes are allowed here) -cm.amq-adapter.route-mapping=otdemo::otdemo:otdemo.localhost.#:5:0 +# cm-dev route +cm.amq-adapter.route-mapping=otdemo:otdemo-exchange:otdemo-queue:api.cm.dev.cleverthis.com.test.python.#:5:0 +# local route +#cm.amq-adapter.route-mapping=otdemo::otdemo:api.dev.localhost.otdemo.#:5:0 # Indicate if AMQ Adapter should respond with HTTP 401 Unauthorized if the request does not contain valid JWT token cm.amq-adapter.require-authenticated-user=false @@ -28,7 +31,7 @@ cm.dispatch.download-dir=/tmp/downloaded_files # Backpressure monitor settings # Define maximum number of messages of being processed in parallel before triggering backpressure OVERLOAD alert -cm.backpressure.threshold=5 +cm.backpressure.threshold=1 # Define maximum CPU usage (%) before triggering backpressure OVERLOAD alert cm.backpressure.threshold-cpu-overload=90 # Define maximum CPU usage (%) before triggering backpressure IDLE alert diff --git a/pyproject.toml b/pyproject.toml index 79de392..87be1c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "amq_adapter" -version = "0.2.33" +version = "0.2.34" description = "The Python implementation of the AMQ Adapter for CleverMicro" readme = "README.md" authors = [ @@ -29,7 +29,8 @@ dependencies = [ "opentelemetry-instrumentation-pika", "pika", "psutil", - "python_multipart" + "python-multipart", + "consul-kv" ] [tool.setuptools.packages.find] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..798a2a6 --- /dev/null +++ b/requirements.txt @@ -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 diff --git a/tests/adapter/perf_test_consul_global_id_generator.py b/tests/adapter/perf_test_consul_global_id_generator.py new file mode 100644 index 0000000..50520b2 --- /dev/null +++ b/tests/adapter/perf_test_consul_global_id_generator.py @@ -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}") diff --git a/tests/adapter/test_backpressure_handler.py b/tests/adapter/test_backpressure_handler.py index 42b7219..a43a7ee 100644 --- a/tests/adapter/test_backpressure_handler.py +++ b/tests/adapter/test_backpressure_handler.py @@ -1,512 +1,205 @@ -# -# NOTE: need to install: pytest-asyncio -# - import asyncio -import json -from asyncio import AbstractEventLoop -from threading import Thread +import time +import unittest from unittest.mock import AsyncMock, MagicMock, patch -import pytest -from aio_pika import Exchange -from aio_pika.abc import AbstractRobustChannel - from amqp.adapter.backpressure_handler import ( BackpressureHandler, - RunningAverage, ScaleRequestV1, + ScalingRequestAlert, ) -from amqp.config.amq_configuration import AMQConfiguration -from amqp.model.model import ScalingRequestAlert -class TestScaleRequestV1: - """ - A class containing pytest unit tests for the ScaleRequestV1 class. - """ +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 - def test_scale_request_v1_initialization(self): - """ - Test that the ScaleRequestV1 object is initialized correctly. - """ - service_id = "test-service" - task_id = "test-task" - max_availability = 10 - current_availability = 5 - request_type = ScalingRequestAlert.UPDATE - - scale_request = ScaleRequestV1( - service_id, - task_id, - max_availability, - current_availability, - request_type, - ) - - assert scale_request.version == 1 - assert scale_request.serviceId == service_id - assert scale_request.taskId == task_id - assert scale_request.maxAvailability == max_availability - assert scale_request.currentAvailability == current_availability - assert scale_request.requestType == request_type - assert scale_request.thread is None - - def test_scale_request_v1_to_json(self): - """ - Test the to_json method of ScaleRequestV1. - """ - service_id = "test-service" - task_id = "test-task" - max_availability = 10 - current_availability = 5 - request_type = ScalingRequestAlert.UPDATE - - scale_request = ScaleRequestV1( - service_id, - task_id, - max_availability, - current_availability, - request_type, - ) - - json_output = scale_request.to_json() - expected_json = json.dumps( - { - "version": 1, - "serviceId": service_id, - "taskId": task_id, - "maxAvailability": max_availability, - "currentAvailability": current_availability, - "requestType": request_type.value, - }, - indent=2, - ) - assert json_output == expected_json - - def test_scale_request_v1_to_json_different_alert_type(self): - """ - Test the to_json method with a different ScalingRequestAlert type. - """ - service_id = "test-service" - task_id = "test-task" - max_availability = 10 - current_availability = 5 - request_type = ScalingRequestAlert.OVERLOAD # Change the alert type - - scale_request = ScaleRequestV1( - service_id, - task_id, - max_availability, - current_availability, - request_type, - ) - - json_output = scale_request.to_json() - expected_json = json.dumps( - { - "version": 1, - "serviceId": service_id, - "taskId": task_id, - "maxAvailability": max_availability, - "currentAvailability": current_availability, - "requestType": request_type.value, # Ensure the value is correctly serialized - }, - indent=2, - ) - assert json_output == expected_json - - -class TestRunningAverage: - """ - A class containing pytest unit tests for the RunningAverage class. - """ - - def test_running_average_initialization(self): - """ - Test that the RunningAverage object is initialized correctly. - """ - num_items = 5 - ra = RunningAverage(num_items) - assert ra.buffer == [0] * num_items - assert ra.pointer == 0 - assert ra.num_items == num_items - assert ra.is_filled is False - - def test_running_average_add_one_value(self): - """ - Test adding a single value to the RunningAverage. - """ - ra = RunningAverage(3) - ra.add(10) - assert ra.buffer == [10, 0, 0] - assert ra.pointer == 1 - assert ra.is_filled is False - - def test_running_average_add_multiple_values_within_capacity(self): - """ - Test adding multiple values within the capacity of the RunningAverage buffer. - """ - ra = RunningAverage(3) - ra.add(5) - ra.add(10) - ra.add(15) - assert ra.buffer == [5, 10, 15] - assert ra.pointer == 0 - assert ra.is_filled is True - - def test_running_average_add_more_values_than_capacity(self): - """ - Test adding more values than the capacity of the RunningAverage buffer, checking wrap-around. - """ - ra = RunningAverage(3) - ra.add(1) - ra.add(2) - ra.add(3) - ra.add(4) # Overwrite the first value - ra.add(5) # Overwrite the second value - assert ra.buffer == [4, 5, 3] - assert ra.pointer == 2 - assert ra.is_filled is True - - def test_running_average_get_average_empty(self): - """ - Test get_average when no values have been added. - """ - ra = RunningAverage(3) - assert ra.get_average() == 0.0 - - def test_running_average_get_average_one_value(self): - """ - Test get_average after adding one value. - """ - ra = RunningAverage(3) - ra.add(10) - assert ra.get_average() == 10.0 - - def test_running_average_get_average_multiple_values_within_capacity(self): - """ - Test get_average with multiple values added within the buffer's capacity. - """ - ra = RunningAverage(3) - ra.add(1) - ra.add(2) - ra.add(3) - assert ra.get_average() == 2.0 - - def test_running_average_get_average_more_values_than_capacity(self): - """ - Test get_average after adding more values than the buffer's capacity (check wrap-around). - """ - ra = RunningAverage(3) - ra.add(1) - ra.add(2) - ra.add(3) - ra.add(4) - ra.add(5) - assert ra.get_average() == 4.0 - - def test_get_current_values_empty(self): - """Test get_current_values when no values have been added.""" - ra = RunningAverage(3) - assert ra.get_current_values() == [] - - def test_get_current_values_one_value(self): - """Test get_current_values after adding one value.""" - ra = RunningAverage(3) - ra.add(10) - assert ra.get_current_values() == [10] - - def test_get_current_values_multiple_values_within_capacity(self): - """Test get_current_values with multiple values within capacity.""" - ra = RunningAverage(3) - ra.add(1) - ra.add(2) - ra.add(3) - assert ra.get_current_values() == [1, 2, 3] - - def test_get_current_values_more_values_than_capacity(self): - """Test get_current_values after adding more values than capacity (wrap-around).""" - ra = RunningAverage(3) - ra.add(1) - ra.add(2) - ra.add(3) - ra.add(4) - ra.add(5) - assert ra.get_current_values() == [3, 4, 5] - - def test_get_current_values_wrap_around(self): - """Test get_current_values when the buffer has wrapped around.""" - ra = RunningAverage(4) - ra.add(1) - ra.add(2) - ra.add(3) - ra.add(4) - ra.add(5) # Overwrite 1 - ra.add(6) # Overwrite 2 - assert ra.get_current_values() == [3, 4, 5, 6] - - -# Test for BackpressureHandler -class TestBackpressureHandler: - """ - A class containing pytest unit tests for the BackpressureHandler class. - """ - - @pytest.fixture - def mock_channel(self): - """ - Pytest fixture to create a mock AbstractRobustChannel. - """ - return AsyncMock(spec=AbstractRobustChannel) - - @pytest.fixture - def mock_loop(self): - """ - Pytest fixture to create a mock AbstractEventLoop. - """ - return MagicMock(spec=AbstractEventLoop) - - @pytest.fixture - def mock_config(self): - """ - Pytest fixture to create a mock AMQConfiguration. - """ - return AMQConfiguration("") - - def test_backpressure_handler_initialization(self, mock_channel, mock_loop, mock_config): - """ - Test that the BackpressureHandler object is initialized correctly. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - assert handler.channel == mock_channel - assert handler.loop == mock_loop - assert handler.config == mock_config - assert handler.swarm_service_id == "clever-amqp-service" - assert handler.swarm_task_id == "python-adapter" - assert handler.parallel_workers == mock_config.backpressure.threshold_threads - assert handler.average_cpu_usage is None - - def test_increase_parallel_executions(self, mock_channel, mock_loop, mock_config): - """ - Test the increase_parallel_executions method. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - initial_executions = handler.current_parallel_executions - handler.increase_parallel_executions() - assert handler.current_parallel_executions == initial_executions + 1 - - def test_decrease_parallel_executions(self, mock_channel, mock_loop, mock_config): - """ - Test the decrease_parallel_executions method. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.current_parallel_executions = 5 - handler.decrease_parallel_executions() - assert handler.current_parallel_executions == 4 - - handler.current_parallel_executions = 0 - handler.decrease_parallel_executions() - assert handler.current_parallel_executions == 0 # Should not go below 0 - - def test_update_parallel_executions(self, mock_channel, mock_loop, mock_config): - """ - Test the update_parallel_executions method. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.update_parallel_executions(100) - assert handler.current_parallel_executions == 100 - - def test_update_last_data_message_time(self, mock_channel, mock_loop, mock_config): - """ - Test the update_last_data_message_time method. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - with patch("time.time") as mock_time: - mock_time.return_value = 12345 - handler.update_last_data_message_time() - assert handler.last_data_message_time == 12345 - - def test_start_backpressure_monitor(self, mock_channel, mock_loop, mock_config): - """ - Test the start_backpressure_monitor method. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.do_loop = 1 # allow only 1 execution of the loop and then exit the thread, else tests will never end - thread = handler.start_backpressure_monitor() - assert isinstance(thread, Thread) - assert handler.thread == thread - - @pytest.mark.asyncio - async def test_check_overload_condition_no_overload(self, mock_channel, mock_loop, mock_config): - """ - Test the check_overload_condition method when there is no overload. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.current_parallel_executions = 1 - handler.parallel_workers = 5 - handler.last_backpressure_event = ScalingRequestAlert.UPDATE - with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload: - await handler.check_overload_condition() - mock_handle_overload.assert_not_called() - - @pytest.mark.asyncio - async def test_check_overload_condition_overload(self, mock_channel, mock_loop, mock_config): - """ - Test the check_overload_condition method when there is an overload. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.current_parallel_executions = 9 # Assuming parallel_workers is 10 - handler.last_backpressure_event = ScalingRequestAlert.UPDATE - with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload: - with patch("time.time") as mock_time: - mock_time.return_value = 12345 - await handler.check_overload_condition() - mock_handle_overload.assert_called_once() - assert handler.last_backpressure_event_time == 12345 - assert handler.last_backpressure_event == ScalingRequestAlert.OVERLOAD - - @pytest.mark.asyncio - async def test_check_overload_condition_already_overload( - self, mock_channel, mock_loop, mock_config - ): - """ - Test the check_overload_condition method when the last event was already an overload. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.current_parallel_executions = 10 - handler.last_backpressure_event = ScalingRequestAlert.OVERLOAD - with patch.object(handler, "handle_backpressure_overload_event") as mock_handle_overload: - await handler.check_overload_condition() - mock_handle_overload.assert_not_called() - - # backpressure_monitor_loop is difficult to test directly without significant mocking - # and potentially long-running tests. The following is a basic outline of how - # you might approach testing it, focusing on verifying that the correct methods - # are called under specific conditions. This is NOT a complete, runnable test. - @pytest.mark.asyncio - async def test_backpressure_monitor_loop(self, mock_channel, mock_loop, mock_config): - """ - Test the backpressure_monitor_loop method (outline). - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.average_cpu_usage = RunningAverage( - 100, initial_value=99 - ) # Mock the RunningAverage in OVERLOAD state - handler.average_cpu_usage.is_filled = True # Together with 0 values ensures OVERLOAD state - # Setup mocks for dependent methods - handler.handle_backpressure_overload_event = AsyncMock() - handler._handle_backpressure_idle_event = AsyncMock() - handler.publish_backpressure_request = AsyncMock() - - # 1. Simulate an overload condition: - handler.last_backpressure_event_time = 0 # Force event trigger - handler.last_data_message_time = 0 - handler.do_loop = 1 - handler.config.backpressure.cpu_overload_duration = ( - -1 - ) # Force overload condition even for recent change - await handler._backpressure_monitor() - handler.handle_backpressure_overload_event.assert_called_once() - - # 2. Simulate an idle condition - handler.average_cpu_usage = RunningAverage( - 100, initial_value=0 - ) # Mock the RunningAverage in IDLE state - handler.average_cpu_usage.is_filled = True # Together with 0 values ensures IDLE state - handler.last_backpressure_event_time = 0 - handler.last_data_message_time = 0 - handler.do_loop = 1 - handler.config.backpressure.idle_duration = ( - -1 - ) # Force idle condition even for recent change - await handler._backpressure_monitor() - handler._handle_backpressure_idle_event.assert_called_once() - - # 3. Simulate neither idle nor overload, but time_window passed - handler.handle_backpressure_overload_event.call_count = 0 - handler._handle_backpressure_idle_event.call_count = 0 - handler.average_cpu_usage = RunningAverage(100) # Mock the RunningAverage in pritine state - handler.last_backpressure_event_time = 0 - handler.last_data_message_time = 0 - handler.do_loop = 1 - await handler._backpressure_monitor() - assert ( - handler.publish_backpressure_request.call_count - + handler.handle_backpressure_overload_event.call_count - + handler._handle_backpressure_idle_event.call_count - ) > 0 - - @pytest.mark.asyncio - async def test_handle_backpressure_overload_event(self, mock_channel, mock_loop, mock_config): - """ - Test the handle_backpressure_overload_event method. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.publish_backpressure_request = AsyncMock() - with patch("time.time") as mock_time: - mock_time.return_value = 12345 - await handler.handle_backpressure_overload_event() - assert handler.publish_backpressure_request.call_count == 1 - assert handler.last_data_message_time == 12345 - - @pytest.mark.asyncio - async def test_handle_backpressure_idle_event(self, mock_channel, mock_loop, mock_config): - """ - Test the _handle_backpressure_idle_event method. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.publish_backpressure_request = AsyncMock() - with patch("time.time") as mock_time: - mock_time.return_value = 12345 - await handler._handle_backpressure_idle_event() - assert handler.publish_backpressure_request.call_count == 1 - assert handler.last_data_message_time == 12345 - - @pytest.mark.asyncio - async def test_publish_backpressure_request_exchange_exists( - self, mock_channel, mock_loop, mock_config - ): - """ - Test the publish_backpressure_request method when the exchange already exists. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.exchange = AsyncMock(spec=Exchange) - with patch("asyncio.run_coroutine_threadsafe") as mock_async_run: - mock_async_run.return_value = asyncio.Future() - mock_async_run.return_value.set_result(None) - await handler.publish_backpressure_request( - ScaleRequestV1("test_service", "test_task", 100, 50, ScalingRequestAlert.UPDATE) - ) - assert mock_async_run.call_count == 1 - # handler.exchange.publish.assert_called_once() - - @pytest.mark.asyncio - async def test_publish_backpressure_request_exchange_does_not_exist( - self, mock_channel, mock_loop, mock_config - ): - """ - Test the publish_backpressure_request method when the exchange does not exist and needs to be created. - """ - handler = BackpressureHandler(mock_channel, mock_loop, mock_config) - handler.exchange = None # Simulate exchange not existing - mock_exchange = AsyncMock(spec=Exchange) - mock_channel.get_exchange.return_value = mock_exchange + # Record callbacks for testing BackpressureHandler._callback_list = {} - with patch("asyncio.run_coroutine_threadsafe") as mock_async_run: - mock_async_run.return_value = asyncio.Future() - mock_async_run.return_value.set_result( - mock_exchange - ) # Make it return the mock exchange - await handler.publish_backpressure_request( - ScaleRequestV1("test_service", "test_task", 100, 50, ScalingRequestAlert.UPDATE) - ) - assert mock_async_run.call_count == 2 + self.handler = BackpressureHandler(self.channel, self.loop, self.config) + self.handler.exchange = AsyncMock() - assert "_wrap_rabbit_mq_api_init" in BackpressureHandler._callback_list - assert "_wrap_rabbit_mq_api" in BackpressureHandler._callback_list - # Since I need also to mockup the asyncio.run_coroutine_threadsafe, the following mockups are - # not called because mockup overrides that call with preset return value. - # mock_channel.get_exchange.assert_called_once_with(name="cleverthis.clevermicro.management") - # handler.exchange.publish.assert_called_once() + # 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() diff --git a/tests/adapter/test_consul_global_id_generator.py b/tests/adapter/test_consul_global_id_generator.py new file mode 100644 index 0000000..acb5087 --- /dev/null +++ b/tests/adapter/test_consul_global_id_generator.py @@ -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() diff --git a/tests/adapter/test_data_message_factory.py b/tests/adapter/test_data_message_factory.py index 9b016ba..202a529 100644 --- a/tests/adapter/test_data_message_factory.py +++ b/tests/adapter/test_data_message_factory.py @@ -107,7 +107,7 @@ class DataMessageFactoryTest(TestCase): self.assertLessEqual( _x_timestamp - _timestamp, 1, - "Timestamp in SnowflakeId differs too much from current time", + "Timestamp in SnowflakeId differs too much from current_availability time", ) def test_to_string(self): diff --git a/tests/config/__init__.py b/tests/config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_amq_configuration.py b/tests/config/test_amq_configuration.py similarity index 100% rename from tests/test_amq_configuration.py rename to tests/config/test_amq_configuration.py diff --git a/tests/rabbitmq/test_rabbit_mq_client.py b/tests/rabbitmq/test_rabbit_mq_client.py index 65da09e..7cfae4e 100644 --- a/tests/rabbitmq/test_rabbit_mq_client.py +++ b/tests/rabbitmq/test_rabbit_mq_client.py @@ -13,10 +13,10 @@ class TestRabbitMQClient(IsolatedAsyncioTestCase): 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.rabbitmq.rabbit_mq_client.RouterConsumer") + patcher = patch("amqp.router.router_consumer.RouterConsumer") self.addCleanup(patcher.stop) self.mock_router = patcher.start() - self.client = RabbitMQClient(self.mock_config) + 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.""" 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..8a4e7fd --- /dev/null +++ b/tests/rabbitmq/test_user_management_service_client.py @@ -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() diff --git a/tests/service/__init__.py b/tests/service/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/service/test_amq_message_handler.py b/tests/service/test_amq_message_handler.py new file mode 100644 index 0000000..f8eda80 --- /dev/null +++ b/tests/service/test_amq_message_handler.py @@ -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() diff --git a/tests/service/test_amq_service.py b/tests/service/test_amq_service.py index 0c57e9c..7f5ff6f 100644 --- a/tests/service/test_amq_service.py +++ b/tests/service/test_amq_service.py @@ -1,185 +1,87 @@ +import asyncio import unittest -from unittest.mock import AsyncMock, MagicMock, Mock, patch +from unittest.mock import AsyncMock, MagicMock, patch -from amqp.config.amq_configuration import AMQAdapter, AMQConfiguration -from amqp.model.model import DataMessage from amqp.service.amq_service import AMQService class TestAMQService(unittest.IsolatedAsyncioTestCase): def setUp(self): - # Create a mock configuration - self.mock_adapter = Mock(spec=AMQAdapter) - self.mock_adapter.generator_id = 12345 - self.mock_adapter.service_name = "test-service" - self.mock_adapter.route_mapping = {} - self.mock_adapter.swarm_service_id = "test-swarm-id" - self.mock_adapter.swarm_task_id = "test-task-id" + # 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 - self.mock_config = Mock(spec=AMQConfiguration) - self.mock_config.amq_adapter = self.mock_adapter + # Mock the service adapter + self.service_adapter = MagicMock() - # Add dispatch configuration - self.mock_dispatch = Mock() - self.mock_dispatch.amq_host = "localhost" - self.mock_dispatch.amq_port = 5672 - self.mock_dispatch.rabbit_mq_user = "guest" - self.mock_dispatch.rabbit_mq_password = "guest" - self.mock_dispatch.use_dlq = False - self.mock_config.dispatch = self.mock_dispatch + # 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) - # Add backpressure configuration - self.mock_backpressure = Mock() - self.mock_backpressure.threshold_threads = 10 - self.mock_config.backpressure = self.mock_backpressure + # 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() - # Create service instance with mocked configuration - self.service = AMQService(self.mock_config) + 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) - # Mock the event loop - self.mock_loop = Mock() - self.service.loop = self.mock_loop + # 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() - def tearDown(self): - # Clean up any resources - if hasattr(self, "service"): - self.service.cleanup() + # The second argument should be the loop + loop_arg = mock_run.call_args[1]["loop"] - @patch("amqp.service.amq_service.RabbitMQClient") - @patch("amqp.service.amq_service.DataMessageFactory") - @patch("amqp.service.amq_service.ServiceMessageFactory") - @patch("amqp.service.amq_service.initialize_telemetry") - def test_init( - self, mock_telemetry, mock_service_factory, mock_data_factory, mock_rabbit_client - ): - # Test initialization of AMQService - service = AMQService(self.mock_config) + self.assertEqual(loop_arg, self.service.backpressure_handler.loop) - # Verify that all required components are initialized - self.assertIsNotNone(service.amq_configuration) - self.assertIsNotNone(service.data_message_factory) - self.assertIsNotNone(service.rabbit_mq_client) - self.assertIsNotNone(service.tracer) - self.assertIsNotNone(service.service_message_factory) + # We can also verify the arguments passed to update_backpressure_value + self.service.backpressure_handler.update_backpressure_value.assert_called_with(50, 200) - @patch("amqp.service.amq_service.RabbitMQClient") - async def test_register_routes_valid(self, mock_rabbit_client): - # Setup mock route mapping - self.mock_adapter.route_mapping = "test-service:test-exchange:test-queue:test.#:5:0" + async def test_backpressure_without_maximum(self): + """Test backpressure method without a maximum value.""" + # Set a maximum_availability value + self.service.maximum_availability = 100 - # Setup mock rabbit client - mock_client = Mock() - mock_client.register_inbound_routes = AsyncMock() - mock_client.register_data_reply_callback = AsyncMock() - mock_client.init = AsyncMock() - mock_client.request_routes = AsyncMock() - mock_client.channel = Mock() - mock_client.channel.is_closed = False - self.service.rabbit_mq_client = mock_client + with patch("asyncio.run_coroutine_threadsafe") as mock_run: + # Call the backpressure method without a maximum value + self.service.backpressure(current_availability=50) - # Setup mock service adapter - mock_service_adapter = Mock() - mock_service_adapter.loop = self.mock_loop + # 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() - # Setup mock message handler - mock_handler = Mock() - mock_handler.inbound_data_message_callback_no_thread = AsyncMock() - mock_handler.reply_received_callback = AsyncMock() - self.service.amq_data_message_handler = mock_handler + 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 - # Initialize the service - await self.service.init(self.mock_loop, mock_service_adapter) + 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) - # Call register_routes - await self.service.register_routes() + # 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() - # Verify that routes were registered twice (once in init, once in register_routes) - self.assertEqual(mock_client.register_inbound_routes.call_count, 2) - self.assertEqual(mock_client.register_data_reply_callback.call_count, 2) + 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 - @patch("amqp.service.amq_service.RabbitMQClient") - async def test_register_routes_invalid(self, mock_rabbit_client): - # Setup invalid route mapping - self.mock_adapter.route_mapping = None + # Call the backpressure method + self.service.backpressure(current_availability=50) - # Setup mock rabbit client - mock_client = Mock() - mock_client.register_inbound_routes = AsyncMock() - mock_client.register_data_reply_callback = AsyncMock() - self.service.rabbit_mq_client = mock_client + # Nothing should happen, no exception should be raised - # Call register_routes - result = await self.service.register_routes() - # Verify that no routes were registered - mock_client.register_inbound_routes.assert_not_called() - mock_client.register_data_reply_callback.assert_not_called() - self.assertIsNone(result) - - @patch("amqp.service.amq_service.RabbitMQClient") - async def test_reinitialize_if_disconnected(self, mock_rabbit_client): - # Setup mock route mapping - self.mock_adapter.route_mapping = "test-service:test-exchange:test-queue:test.#:5:0" - - # Setup mock rabbit client with closed channel - mock_client = AsyncMock() - mock_client.channel = Mock() - mock_client.channel.is_closed = True - mock_client.init = AsyncMock() - mock_client.register_inbound_routes = AsyncMock() - mock_client.register_data_reply_callback = AsyncMock() - mock_client.request_routes = AsyncMock() - - # Make the constructor return our mock client - mock_rabbit_client.return_value = mock_client - self.service.rabbit_mq_client = mock_client - - # Setup mock message handler - mock_handler = Mock() - mock_handler.inbound_data_message_callback_no_thread = AsyncMock() - mock_handler.reply_received_callback = AsyncMock() - self.service.amq_data_message_handler = mock_handler - - # Call reinitialize_if_disconnected - await self.service.reinitialize_if_disconnected() - - # Verify that client was reinitialized - mock_client.init.assert_called_once() - - def test_send_message(self): - # Create mock message and future - mock_message = Mock(spec=DataMessage) - mock_message.id.return_value = "test-id" - mock_future = MagicMock() - - # Setup mock message handler - self.service.amq_data_message_handler = Mock() - self.service.amq_data_message_handler.outstanding = {} - - # Send message - self.service.send_message(mock_message, mock_future) - - # Verify message was added to outstanding - self.assertEqual(self.service.amq_data_message_handler.outstanding["test-id"], mock_future) - - def test_remove_outstanding_future(self): - # Setup mock message handler with outstanding message - self.service.amq_data_message_handler = Mock() - self.service.amq_data_message_handler.outstanding = {"test-id": Mock()} - - # Remove future - self.service.remove_outstanding_future("test-id") - - # Verify message was removed from outstanding - self.assertNotIn("test-id", self.service.amq_data_message_handler.outstanding) - - def test_cleanup(self): - # Setup mock rabbit client - mock_client = Mock() - self.service.rabbit_mq_client = mock_client - - # Call cleanup - self.service.cleanup() - - # Verify rabbit client cleanup was called - mock_client.cleanup.assert_called_once() +if __name__ == "__main__": + unittest.main()