From: Thorsten Date: Sat, 18 Jul 2026 07:16:02 +0000 (+0200) Subject: Add NATS/JetStream as an alternative broker to RabbitMQ X-Git-Url: https://git.aero2k.de/?a=commitdiff_plain;h=b82ee8602c4c262f00d8373e236ac201b0bdf6a8;p=urlbot-v3.git Add NATS/JetStream as an alternative broker to RabbitMQ RabbitMQ's Erlang VM baseline costs ~300-500MB RES regardless of actual traffic, which is wasteful for a small hobby-scale chatbot. Introduces a runtime toggle (broker_backend = amqp | nats in local_config.ini) so a deployment can opt into a NATS+JetStream backend instead, while defaulting to amqp so existing deployments are unaffected. - New common/broker.py abstraction (SyncBroker/AsyncBroker) with broker_amqp.py (thin, behavior-preserving wrapper around the existing pika calls) and broker_nats.py backends. - Binding-key wildcard matching (some patterns, e.g. fun.py's mid-pattern "#", have no native NATS subject-algebra equivalent) is handled by subscribing broadly and filtering in-process via common/routing.py, lifted from the AMQP-routing simulator already used in integration tests. - action_processing and plugin_registry go through JetStream for durability/late-subscriber buffering, matching today's AMQP guarantees. - Ansible role: broker_backend/nats_uri toggle, optional local NATS server provisioning (nats_local), wait-for-broker.sh picked conditionally. Co-Authored-By: Claude Sonnet 5 --- diff --git a/deploy/roles/urlbot/defaults/main.yml b/deploy/roles/urlbot/defaults/main.yml index ecd8315..8e7351f 100644 --- a/deploy/roles/urlbot/defaults/main.yml +++ b/deploy/roles/urlbot/defaults/main.yml @@ -1,3 +1,8 @@ --- # local user to be used -botuser: jabberbot \ No newline at end of file +botuser: jabberbot +# message broker backend: "amqp" (RabbitMQ, default) or "nats" (NATS + JetStream) +broker_backend: amqp +# whether to install and run a local NATS server on this host (see tasks/nats_server.yml) +nats_local: false +nats_version: "2.14.3" \ No newline at end of file diff --git a/deploy/roles/urlbot/handlers/main.yml b/deploy/roles/urlbot/handlers/main.yml index 4b70dc7..02874a7 100644 --- a/deploy/roles/urlbot/handlers/main.yml +++ b/deploy/roles/urlbot/handlers/main.yml @@ -10,3 +10,9 @@ scope: user name: urlbot-worker.service state: restarted + +- name: restart nats-server + become: true + systemd: + name: nats-server.service + state: restarted diff --git a/deploy/roles/urlbot/tasks/main.yml b/deploy/roles/urlbot/tasks/main.yml index fdfd062..a172846 100644 --- a/deploy/roles/urlbot/tasks/main.yml +++ b/deploy/roles/urlbot/tasks/main.yml @@ -1,5 +1,9 @@ --- +- name: provision local nats server + include_tasks: nats_server.yml + when: broker_backend == 'nats' and nats_local | default(false) + - name: setup virtualenv for the chatbot shell: cmd: "python3 -m venv {{ venv_chatbot }}" @@ -47,7 +51,7 @@ - key: "password" value: "{{password}}" - key: "rooms" - value: "{{rooms | join(', ')}}" + value: "{{rooms | join(', ')}}," - key: "src-url" value: "{{botrepo_view}}" - key: "bot_nickname" @@ -58,6 +62,8 @@ value: "{{bot_owner_email}}" - key: "amqp_uri" value: "{{ amqp_uri }}" + - key: "broker_backend" + value: "{{ broker_backend }}" # TODO: detectlanguage_api_key, sudoers and giphy_key - key: "giphy_key" value: "{{giphy_key}}" @@ -69,13 +75,21 @@ - restart worker - restart chatbot +- name: set nats configuration + lineinfile: dest=~/urlbot-v3/local_config.ini create=yes line="nats_uri = {{ nats_uri }}" regexp="^nats_uri.=" + when: broker_backend == 'nats' + tags: [render_config] + notify: + - restart worker + - restart chatbot + - name: create directory for systemd unit files shell: mkdir -p ~/.config/systemd/user/ creates=~/.config/systemd/user/ - name: unitfile support template: - src: "wait-for-rabbitmq.sh" - dest: "~/wait-for-rabbitmq.sh" + src: "{{ 'wait-for-nats.sh' if broker_backend == 'nats' else 'wait-for-rabbitmq.sh' }}" + dest: "~/wait-for-broker.sh" mode: u+x - name: unitfile diff --git a/deploy/roles/urlbot/tasks/nats_server.yml b/deploy/roles/urlbot/tasks/nats_server.yml new file mode 100644 index 0000000..247e844 --- /dev/null +++ b/deploy/roles/urlbot/tasks/nats_server.yml @@ -0,0 +1,72 @@ +--- +# Installs and runs a local NATS+JetStream server on this host. Only +# included when broker_backend == 'nats' and nats_local is true - opt-in per +# host, unlike RabbitMQ which this role has always expected to be +# pre-provisioned externally. + +- name: nats - create system user + become: true + user: + name: nats + system: true + shell: /usr/sbin/nologin + create_home: false + +- name: nats - create JetStream storage directory + become: true + file: + path: /var/lib/nats/jetstream + state: directory + owner: nats + group: nats + mode: "0750" + +- name: nats - determine release architecture + set_fact: + nats_release_arch: "{{ { + 'x86_64': 'amd64', + 'aarch64': 'arm64', + 'armv7l': 'arm7', + 'armv6l': 'arm6', + }[ansible_architecture] }}" + +- name: nats - download server release + become: true + get_url: + url: "https://github.com/nats-io/nats-server/releases/download/v{{ nats_version }}/nats-server-v{{ nats_version }}-linux-{{ nats_release_arch }}.tar.gz" + dest: "/tmp/nats-server-v{{ nats_version }}-linux-{{ nats_release_arch }}.tar.gz" + register: nats_download + +- name: nats - unpack and install binary + become: true + unarchive: + src: "/tmp/nats-server-v{{ nats_version }}-linux-{{ nats_release_arch }}.tar.gz" + dest: /tmp + remote_src: true + when: nats_download is changed + +- name: nats - copy binary into place + become: true + copy: + src: "/tmp/nats-server-v{{ nats_version }}-linux-{{ nats_release_arch }}/nats-server" + dest: /usr/local/bin/nats-server + mode: "0755" + remote_src: true + when: nats_download is changed + notify: + - restart nats-server + +- name: nats - install systemd unit + become: true + template: + src: nats-server.service + dest: /etc/systemd/system/nats-server.service + register: nats_unit + +- name: nats - enable and start service + become: true + systemd: + name: nats-server + daemon_reload: true + enabled: true + state: started diff --git a/deploy/roles/urlbot/templates/nats-server.service b/deploy/roles/urlbot/templates/nats-server.service new file mode 100644 index 0000000..4ef15c2 --- /dev/null +++ b/deploy/roles/urlbot/templates/nats-server.service @@ -0,0 +1,12 @@ +[Unit] +Description=NATS server (JetStream enabled) +After=network.target + +[Service] +ExecStart=/usr/local/bin/nats-server -js -sd /var/lib/nats/jetstream +User=nats +Group=nats +Restart=always + +[Install] +WantedBy=multi-user.target diff --git a/deploy/roles/urlbot/templates/urlbot-chat.service b/deploy/roles/urlbot/templates/urlbot-chat.service index b2eaef9..604d0bb 100644 --- a/deploy/roles/urlbot/templates/urlbot-chat.service +++ b/deploy/roles/urlbot/templates/urlbot-chat.service @@ -2,7 +2,7 @@ Description=jabber bot entertaining and supporting activity on jabber MUCs [Service] -ExecStartPre=/home/{{ botuser }}/wait-for-rabbitmq.sh +ExecStartPre=/home/{{ botuser }}/wait-for-broker.sh ExecStart={{ venv_chatbot }}/bin/urlbotd-chat WorkingDirectory=/home/{{ botuser }}/urlbot-v3/ StandardOutput=journal+console diff --git a/deploy/roles/urlbot/templates/urlbot-worker.service b/deploy/roles/urlbot/templates/urlbot-worker.service index a035685..94b03e1 100644 --- a/deploy/roles/urlbot/templates/urlbot-worker.service +++ b/deploy/roles/urlbot/templates/urlbot-worker.service @@ -2,7 +2,7 @@ Description=jabber bot entertaining and supporting activity on jabber MUCs [Service] -ExecStartPre=/home/{{ botuser }}/wait-for-rabbitmq.sh +ExecStartPre=/home/{{ botuser }}/wait-for-broker.sh ExecStart={{ venv_worker }}/bin/urlbotd-worker WorkingDirectory=/home/{{ botuser }}/urlbot-v3/ StandardOutput=journal+console diff --git a/deploy/roles/urlbot/templates/wait-for-nats.sh b/deploy/roles/urlbot/templates/wait-for-nats.sh new file mode 100644 index 0000000..37c9fc3 --- /dev/null +++ b/deploy/roles/urlbot/templates/wait-for-nats.sh @@ -0,0 +1,6 @@ +#!/bin/bash + +while ! echo > /dev/tcp/{{ nats_uri | urlsplit('hostname') }}/{{ nats_uri | urlsplit('port') or 4222 }}; do + echo "waiting for remote service..." + sleep 5s +done diff --git a/pyproject.toml b/pyproject.toml index fe03401..3ef24aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ classifiers = [ dependencies = [ "slixmpp>=1.10.0", "pika", + "nats-py", "configobj", "requests", ] @@ -24,6 +25,7 @@ dependencies = [ [project.optional-dependencies] chatbot = [] worker = ["lxml", "plyvel"] +ai = ["torch", "transformers"] [project.scripts] urlbotd-chat = "distbot.bot.bot:run" diff --git a/src/distbot/bot/action_worker.py b/src/distbot/bot/action_worker.py index 14d91eb..02bffa4 100644 --- a/src/distbot/bot/action_worker.py +++ b/src/distbot/bot/action_worker.py @@ -1,16 +1,13 @@ # -*- coding: utf-8 -*- import asyncio -import functools import json import logging import time from asyncio import TimerHandle from typing import List -import pika -from pika.channel import Channel - -from distbot.common.action import Action, send_action +from distbot.common.action import Action +from distbot.common.broker import get_async_broker from distbot.common.config import conf_set, conf_get from distbot.common.message import process_message @@ -29,51 +26,21 @@ class ActionWorker(object): def __init__(self, bot, queue_name): self._bot = bot - self._connection = None + self._broker = get_async_broker() self._channel = None self._queue_name = queue_name - self._consumer_tag = None self.loop = asyncio.get_event_loop() self.idle_task = self.loop.create_task(self.tick()) - def open_connection(self, connection): - logger.debug("opened connection for actionworker") - self._connection = connection - connection.channel(on_open_callback=self.open_channel) - - def open_channel(self, channel: Channel): - logger.debug("opened channel") - # self.add_on_channel_close_callback() - # self.setup_exchange(self.EXCHANGE) - self._channel = channel - channel.add_on_close_callback(lambda *_: self._bot.disconnect()) - channel.queue_declare(queue=self._queue_name, durable=True, callback=lambda _: self.start_consuming()) - def cancel(_): - logger.debug("cancel channel") - channel.close() - channel.add_on_cancel_callback(lambda _: cancel) - - def start_consuming(self): - """This method sets up the consumer by first calling - add_on_cancel_callback so that the object is notified if RabbitMQ - cancels the consumer. It then issues the Basic.Consume RPC command - which returns the consumer tag that is used to uniquely identify the - consumer with RabbitMQ. We keep the value to use it when we want to - cancel consuming. The on_message method is passed in as a callback pika - will invoke when a message is fully received. - - """ - logger.info('Issuing consumer related RPC commands') - # self.add_on_cancel_callback() - self._consumer_tag = self._channel.basic_consume( - queue=self._queue_name, - on_message_callback=self.on_message + async def start(self): + logger.debug("connecting actionworker to durable queue %s", self._queue_name) + self._channel = await self._broker.connect_durable_queue( + self._queue_name, self.on_message, on_channel_closed=self._bot.disconnect, ) - logger.info("Setup consumer for actions on ch %s on con %s with tag %s", self._channel, self._connection, - self._consumer_tag) + logger.info("actionworker consuming from %s", self._queue_name) - def on_message(self, ch, method, properties: pika.spec.BasicProperties, body): + def on_message(self, ch, method, properties, body): logger.info('Received message # %s from %s: %s', method.delivery_tag, properties.app_id, body) body = json.loads(body.decode("utf-8")) @@ -95,13 +62,20 @@ class ActionWorker(object): break return action_item + def _publish_action(self, action: Action): + self._channel.basic_publish( + exchange='', + routing_key=self._queue_name, + body=action.serialize() + ) + def schedule_action(self, action: Action): logger.info("scheduling event: %s", action.serialize()) if action.mutex and self.find_scheduled_action_by_mutex(action.mutex): logger.info("not scheduling that event (prevented by mutex)") raise RuntimeError("not scheduling that event (prevented by mutex)") - handle: TimerHandle = self.loop.call_later(action.time - time.time(), functools.partial(send_action, self._queue_name), action) + handle: TimerHandle = self.loop.call_later(action.time - time.time(), self._publish_action, action) self.event_list.append(handle) @staticmethod @@ -171,5 +145,5 @@ class ActionWorker(object): conf_set('request_counter', request_counter + 1) def die(self): - if self._connection: - self._connection.close() + if self._channel: + asyncio.ensure_future(self._broker.close()) diff --git a/src/distbot/bot/bot.py b/src/distbot/bot/bot.py index a4dcb04..e29f1ab 100644 --- a/src/distbot/bot/bot.py +++ b/src/distbot/bot/bot.py @@ -5,9 +5,6 @@ import logging import re import shlex -import pika -from pika.adapters.asyncio_connection import AsyncioConnection - import slixmpp from slixmpp.exceptions import IqError, IqTimeout from slixmpp.stanza import Message @@ -63,13 +60,7 @@ class Bot(slixmpp.ClientXMPP): self.action_worker.die() def initialize_actionworker(self): - connection = AsyncioConnection( - pika.URLParameters(conf_get("amqp_uri")), - on_open_callback=self.action_worker.open_connection, - on_open_error_callback=lambda con, err: logger.error("Could not connect: %s", err), - on_close_callback=lambda con, e, x=None: logger.info("closing connection of actionworker (%s)", e) - ) - logger.debug("connection state: %s", connection.connection_state) + asyncio.ensure_future(self.action_worker.start()) def disconnect(self, reconnect=False, wait=None, send_close=True): logger.info("Stopping all workers...") diff --git a/src/distbot/bot/worker.py b/src/distbot/bot/worker.py index 4e348d1..813f5b3 100644 --- a/src/distbot/bot/worker.py +++ b/src/distbot/bot/worker.py @@ -6,11 +6,8 @@ from collections import deque, defaultdict from functools import partial from typing import Optional -import pika -import pika.channel -import pika.exceptions - from distbot.common.action import Action, send_action +from distbot.common.broker import SyncBroker, get_sync_broker from distbot.common.config import conf_get logger = logging.getLogger(__name__) @@ -50,32 +47,20 @@ class Worker(threading.Thread): else: self.usage = "(reaction only)" self.used_routing_key = None - self.connection: Optional[pika.BlockingConnection] = None - self.channel = None + self.broker: Optional[SyncBroker] = None def init_channel(self): - self.connection = pika.BlockingConnection(pika.URLParameters(conf_get("amqp_uri"))) - self.channel = self.connection.channel() - - self.channel.exchange_declare(exchange='classifier', exchange_type='topic') + self.broker = get_sync_broker() + self.broker.connect() def init_queue(self): - result = self.channel.queue_declare(exclusive=True, queue="") - self.queue = result.method.queue - - for binding_key in self.binding_keys: - logger.info("Registering plugin %s for %s to queue %s", self.get_subclass_name(), binding_key, self.queue) - self.channel.queue_bind( - exchange='classifier', - queue=self.queue, - routing_key=binding_key - ) + self.broker.register_classifier_consumer(self.binding_keys, self.callback) def callback( self, - ch: pika.channel.Channel, - method: pika.spec.Basic.Deliver, - properties: pika.spec.BasicProperties, + ch, + method, + properties, body: bytes, ): logger.debug("Reacting on %s in %s", str(method.routing_key), self.get_subclass_name()) @@ -111,25 +96,24 @@ class Worker(threading.Thread): return self.__class__.__name__ def die(self): - self.connection.add_callback_threadsafe(callback=self.channel.stop_consuming) + self.broker.stop() def run(self): self.init_channel() self.register_plugin() self.init_queue() - self.channel.basic_consume(queue=self.queue, on_message_callback=self.callback) - self.channel.start_consuming() + self.broker.run_forever() def parse_body(self, msg: dict) -> Action | None: raise NotImplementedError() def register_plugin(self): - self.channel.queue_declare(queue='plugin_registry') - self.channel.basic_publish( - exchange='', - routing_key='plugin_registry', - body=json.dumps(self.get_declaration()).encode("utf-8") + self.broker.publish_queue( + 'plugin_registry', + json.dumps(self.get_declaration()).encode("utf-8"), + durable=False, + max_age_seconds=120, ) def get_declaration(self): diff --git a/src/distbot/common/action.py b/src/distbot/common/action.py index 5b5f727..1ff4224 100644 --- a/src/distbot/common/action.py +++ b/src/distbot/common/action.py @@ -2,8 +2,7 @@ import json from copy import deepcopy -import pika -from distbot.common.config import conf_get +from distbot.common.broker import get_sync_broker class Action: @@ -69,14 +68,9 @@ class Action: def send_action(actionqueue, action: Action): - connection = pika.BlockingConnection( - pika.URLParameters(conf_get("amqp_uri")) - ) - channel = connection.channel() - channel.queue_declare(queue=actionqueue, durable=True) - - channel.basic_publish( - exchange='', - routing_key=actionqueue, - body=action.serialize() - ) + broker = get_sync_broker() + broker.connect() + try: + broker.publish_queue(actionqueue, action.serialize()) + finally: + broker.close() diff --git a/src/distbot/common/broker.py b/src/distbot/common/broker.py new file mode 100644 index 0000000..f319a31 --- /dev/null +++ b/src/distbot/common/broker.py @@ -0,0 +1,98 @@ +# -*- coding: utf-8 -*- +"""Broker abstraction: lets the AMQP (RabbitMQ) and NATS backends be selected +at runtime via the `broker_backend` config key, without call sites caring +which one is in use. + +Two interfaces are needed because the codebase already has two distinct +connection models: + +- `SyncBroker`: blocking, used by one instance per owning thread (a plugin + `Worker` thread, or the process-wide singleton used by common/message.py + and common/action.py). +- `AsyncBroker`: native asyncio, used only by action_worker.py, which already + runs on bot.py's slixmpp event loop. + +Callback signatures deliberately mirror pika's own callback shape +(`on_message(ch, method, properties, body)`, with `method.routing_key`, +`method.delivery_tag`, and `ch.basic_ack(delivery_tag=...)`) so that +`Worker.callback()` and `ActionWorker.on_message()` do not need to change at +all when the backend changes - only connection/consumer setup differs. +""" +from abc import ABC, abstractmethod +from typing import Callable, List, Protocol + +from distbot.common.config import conf_get + +OnMessage = Callable[..., None] + + +class SyncBroker(ABC): + @abstractmethod + def connect(self) -> None: ... + + @abstractmethod + def close(self) -> None: ... + + @abstractmethod + def publish(self, subject: str, body: bytes) -> None: + """Fire-and-forget onto the shared classifier fanout.""" + + @abstractmethod + def publish_queue(self, queue: str, body: bytes, durable: bool = True, max_age_seconds: int = None) -> None: + """Point-to-point publish directly onto a named queue. + + max_age_seconds only matters for backends with a durability layer + that needs an explicit retention window (JetStream); ignored + otherwise (a plain AMQP queue keeps messages until consumed either + way).""" + + @abstractmethod + def register_classifier_consumer(self, binding_keys: List[str], on_message: OnMessage) -> None: + """Register (does not block) a handler for classifier messages whose + subject matches any of binding_keys. Multiple calls accumulate.""" + + @abstractmethod + def register_queue_consumer(self, queue: str, on_message: OnMessage, max_age_seconds: int = None) -> None: + """Register a plain point-to-point queue consumer.""" + + @abstractmethod + def run_forever(self) -> None: + """Blocks the calling thread, dispatching to all registered consumers + until stop() is called (from another thread).""" + + @abstractmethod + def stop(self) -> None: + """Thread-safe; unblocks run_forever().""" + + +class DurableChannel(Protocol): + def basic_publish(self, exchange: str, routing_key: str, body: bytes) -> None: ... + + +class AsyncBroker(ABC): + @abstractmethod + async def connect_durable_queue( + self, queue: str, on_message: OnMessage, on_channel_closed: Callable[[], None] = None, + ) -> DurableChannel: + """Connects, ensures the durable queue exists, starts consuming, and + returns a duck-typed channel exposing basic_publish(exchange, + routing_key, body) for fire-and-forget publishes onto that queue.""" + + @abstractmethod + async def close(self) -> None: ... + + +def get_sync_broker() -> SyncBroker: + if conf_get("broker_backend") == "nats": + from distbot.common.broker_nats import NatsSyncBroker + return NatsSyncBroker() + from distbot.common.broker_amqp import AmqpSyncBroker + return AmqpSyncBroker() + + +def get_async_broker() -> AsyncBroker: + if conf_get("broker_backend") == "nats": + from distbot.common.broker_nats import NatsAsyncBroker + return NatsAsyncBroker() + from distbot.common.broker_amqp import AmqpAsyncBroker + return AmqpAsyncBroker() diff --git a/src/distbot/common/broker_amqp.py b/src/distbot/common/broker_amqp.py new file mode 100644 index 0000000..4486b87 --- /dev/null +++ b/src/distbot/common/broker_amqp.py @@ -0,0 +1,110 @@ +# -*- coding: utf-8 -*- +import asyncio +import logging +from typing import Callable, List + +import pika +from pika.adapters.asyncio_connection import AsyncioConnection + +from distbot.common.broker import AsyncBroker, DurableChannel, OnMessage, SyncBroker +from distbot.common.config import conf_get + +logger = logging.getLogger(__name__) + + +class AmqpSyncBroker(SyncBroker): + """Thin wrapper around the exact pika calls the codebase used directly + before the broker abstraction existed - a behavior-preserving refactor.""" + + def __init__(self): + self.connection = None + self.channel = None + + def connect(self) -> None: + self.connection = pika.BlockingConnection(pika.URLParameters(conf_get("amqp_uri"))) + self.channel = self.connection.channel() + self.channel.exchange_declare(exchange='classifier', exchange_type='topic') + + def close(self) -> None: + if self.connection: + self.connection.close() + + def publish(self, subject: str, body: bytes) -> None: + self.channel.basic_publish(exchange='classifier', routing_key=subject, body=body) + + def publish_queue(self, queue: str, body: bytes, durable: bool = True, max_age_seconds: int = None) -> None: + # max_age_seconds is a JetStream-only retention knob; AMQP has no + # equivalent and keeps messages until consumed regardless. + self.channel.queue_declare(queue=queue, durable=durable) + self.channel.basic_publish(exchange='', routing_key=queue, body=body) + + def register_classifier_consumer(self, binding_keys: List[str], on_message: OnMessage) -> None: + result = self.channel.queue_declare(exclusive=True, queue="") + queue = result.method.queue + for binding_key in binding_keys: + logger.info("Registering consumer for %s on queue %s", binding_key, queue) + self.channel.queue_bind(exchange='classifier', queue=queue, routing_key=binding_key) + self.channel.basic_consume(queue=queue, on_message_callback=on_message) + + def register_queue_consumer(self, queue: str, on_message: OnMessage, max_age_seconds: int = None) -> None: + self.channel.queue_declare(queue=queue) + self.channel.basic_consume(queue=queue, on_message_callback=on_message) + + def run_forever(self) -> None: + self.channel.start_consuming() + + def stop(self) -> None: + self.connection.add_callback_threadsafe(callback=self.channel.stop_consuming) + + +class AmqpAsyncBroker(AsyncBroker): + def __init__(self): + self._connection = None + self._channel = None + + async def connect_durable_queue( + self, queue: str, on_message: OnMessage, on_channel_closed: Callable[[], None] = None, + ) -> DurableChannel: + loop = asyncio.get_event_loop() + connected: asyncio.Future = loop.create_future() + + def _on_open(connection): + logger.debug("opened connection for actionworker") + connection.channel(on_open_callback=_on_channel_open) + + def _on_channel_open(channel): + logger.debug("opened channel") + self._channel = channel + if on_channel_closed: + channel.add_on_close_callback(lambda *_: on_channel_closed()) + + def _cancel(_): + logger.debug("cancel channel") + channel.close() + channel.add_on_cancel_callback(lambda _: _cancel) + + def _queue_declared(_): + channel.basic_consume(queue=queue, on_message_callback=on_message) + if not connected.done(): + connected.set_result(channel) + channel.queue_declare(queue=queue, durable=True, callback=_queue_declared) + + def _on_open_error(_conn, err): + logger.error("Could not connect: %s", err) + if not connected.done(): + connected.set_exception(RuntimeError(str(err))) + + def _on_close(_conn, exc, _x=None): + logger.info("closing connection of actionworker (%s)", exc) + + self._connection = AsyncioConnection( + pika.URLParameters(conf_get("amqp_uri")), + on_open_callback=_on_open, + on_open_error_callback=_on_open_error, + on_close_callback=_on_close, + ) + return await connected + + async def close(self) -> None: + if self._connection: + self._connection.close() diff --git a/src/distbot/common/broker_nats.py b/src/distbot/common/broker_nats.py new file mode 100644 index 0000000..dfadc95 --- /dev/null +++ b/src/distbot/common/broker_nats.py @@ -0,0 +1,243 @@ +# -*- coding: utf-8 -*- +"""NATS/JetStream implementation of the broker abstraction (common/broker.py). + +Design notes (see the migration plan for the full rationale): + +- Binding-key patterns are AMQP topic-exchange syntax baked into ~20 plugins + (some, e.g. fun.py's "me.#..#", have no native NATS-subject-algebra + equivalent - '#' can appear mid-pattern and matches zero-or-more words, + while NATS '>' can only be a trailing token and matches one-or-more). So + the classifier fanout does NOT attempt subject-algebra translation: every + subscriber subscribes to the single broad subject "classifier.>" and + filters in-process with distbot.common.routing.matches(), the same regex + matcher already used to simulate AMQP routing in the test harness. + +- nats-py has no blocking/sync client, but worker.py's thread-per-plugin + model doesn't need to change: each NatsSyncBroker instance runs its own + asyncio event loop on a dedicated background thread and marshals every + call onto it, so the plugin thread itself stays a plain blocking caller. + +- Named point-to-point queues (action_processing, plugin_registry) always + go through JetStream, not core pub/sub, because core NATS has no + buffering for a subscriber that hasn't attached yet - unlike AMQP's real + (if possibly non-durable) queues. plugin_registry gets a short max_age so + stale registrations from restarted/removed plugins don't linger. +""" +import asyncio +import logging +import threading +from typing import Callable, List + +import nats +import nats.js.errors + +from distbot.common.broker import AsyncBroker, DurableChannel, OnMessage, SyncBroker +from distbot.common.config import conf_get +from distbot.common.routing import matches + +logger = logging.getLogger(__name__) + +CLASSIFIER_PREFIX = "classifier." + + +def _to_str(value) -> str: + return value.decode("utf-8") if isinstance(value, bytes) else value + + +def _to_bytes(value) -> bytes: + return value.encode("utf-8") if isinstance(value, str) else value + + +def _stream_name_for(queue: str) -> str: + return queue.upper() + + +def _durable_name_for(queue: str) -> str: + return f"{queue}-consumer" + + +class _Method: + def __init__(self, routing_key, delivery_tag=None): + self.routing_key = routing_key + self.delivery_tag = delivery_tag + + +class _Properties: + app_id = None + + +class _NoAckChannel: + """Used for classifier-fanout dispatch: core NATS pub/sub has no ack + concept, matching AMQP's fanout-to-exclusive-queue semantics today.""" + + def basic_publish(self, exchange: str, routing_key: str, body: bytes) -> None: + raise NotImplementedError("classifier dispatch channel is receive-only") + + def basic_ack(self, delivery_tag=None) -> None: + pass + + +class _JetStreamAckChannel: + def __init__(self, msg): + self._msg = msg + + def basic_ack(self, delivery_tag=None) -> None: + asyncio.ensure_future(self._msg.ack()) + + +async def _ensure_stream(js, queue: str, max_age_seconds: int = None): + stream_name = _stream_name_for(queue) + config = {"name": stream_name, "subjects": [queue]} + if max_age_seconds: + config["max_age"] = max_age_seconds + try: + await js.stream_info(stream_name) + except nats.js.errors.NotFoundError: + await js.add_stream(**config) + + +class NatsSyncBroker(SyncBroker): + """One instance per owning thread. Used both for one-shot connect/publish/ + close call sites (common/message.py, common/action.py) and for the + long-running per-plugin-thread consumer model (worker.py). + + nats-py is asyncio-only, but callers of this class are plain sync code - + including, critically, callbacks invoked *from inside* another + NatsSyncBroker's own dispatch (Worker.callback() calling send_action(), + which builds a brand new NatsSyncBroker of its own). A naive + "new_event_loop() + run_until_complete()" per call breaks the moment the + calling thread already has a loop running (asyncio forbids nesting), so + instead every instance runs its own persistent event loop on a dedicated + background thread, and every operation is marshalled onto it via + run_coroutine_threadsafe(...).result() - safe to call from any thread, + async or not, since it never touches the calling thread's own loop.""" + + def __init__(self): + self._loop = None + self._thread = None + self._nc = None + self._js = None + self._classifier_consumers = [] + self._queue_consumers = [] + self._subs = [] + self._stop_wait_event = None + + def connect(self) -> None: + self._loop = asyncio.new_event_loop() + self._thread = threading.Thread(target=self._loop.run_forever, daemon=True) + self._thread.start() + self._call(self._connect()) + + def _call(self, coro): + return asyncio.run_coroutine_threadsafe(coro, self._loop).result() + + async def _connect(self): + self._nc = await nats.connect(conf_get("nats_uri")) + self._js = self._nc.jetstream() + + def close(self) -> None: + if self._nc is not None: + self._call(self._nc.close()) + if self._loop is not None: + self._loop.call_soon_threadsafe(self._loop.stop) + self._thread.join(timeout=5) + self._loop.close() + + def publish(self, subject: str, body: bytes) -> None: + full_subject = CLASSIFIER_PREFIX + _to_str(subject) + self._call(self._nc.publish(full_subject, _to_bytes(body))) + + def publish_queue(self, queue: str, body: bytes, durable: bool = True, max_age_seconds: int = None) -> None: + self._call(self._publish_queue(queue, body, max_age_seconds)) + + async def _publish_queue(self, queue: str, body: bytes, max_age_seconds: int): + await _ensure_stream(self._js, queue, max_age_seconds) + await self._js.publish(queue, _to_bytes(body)) + + def register_classifier_consumer(self, binding_keys: List[str], on_message: OnMessage) -> None: + self._classifier_consumers.append((binding_keys, on_message)) + + def register_queue_consumer(self, queue: str, on_message: OnMessage, max_age_seconds: int = None) -> None: + self._queue_consumers.append((queue, on_message, max_age_seconds)) + + def run_forever(self) -> None: + self._call(self._start_consumers()) + self._stop_wait_event = threading.Event() + self._stop_wait_event.wait() + + async def _start_consumers(self): + if self._classifier_consumers: + sub = await self._nc.subscribe(CLASSIFIER_PREFIX + ">", cb=self._on_classifier_message) + self._subs.append(sub) + + for queue, handler, max_age_seconds in self._queue_consumers: + await _ensure_stream(self._js, queue, max_age_seconds) + sub = await self._js.subscribe( + queue, durable=_durable_name_for(queue), cb=self._make_queue_callback(handler), manual_ack=True, + ) + self._subs.append(sub) + + async def _on_classifier_message(self, msg): + subject = msg.subject[len(CLASSIFIER_PREFIX):] + method = _Method(routing_key=subject) + for binding_keys, handler in self._classifier_consumers: + if any(matches(binding_key, subject) for binding_key in binding_keys): + handler(_NoAckChannel(), method, _Properties(), msg.data) + break + + def stop(self) -> None: + if self._stop_wait_event is not None: + self._stop_wait_event.set() + + def _make_queue_callback(self, handler: OnMessage): + async def _cb(msg): + method = _Method(routing_key=msg.subject) + handler(_JetStreamAckChannel(msg), method, _Properties(), msg.data) + return _cb + + +class _NatsDurableChannel: + """Duck-typed as common/broker.py's DurableChannel: exposes + basic_publish(exchange, routing_key, body) so action_worker.py's + _publish_action() doesn't need to know which backend it's talking to.""" + + def __init__(self, js, queue: str): + self._js = js + self._queue = queue + + def basic_publish(self, exchange: str, routing_key: str, body: bytes) -> None: + asyncio.ensure_future(self._js.publish(self._queue, _to_bytes(body))) + + +class NatsAsyncBroker(AsyncBroker): + def __init__(self): + self._nc = None + self._js = None + self._sub = None + + async def connect_durable_queue( + self, queue: str, on_message: OnMessage, on_channel_closed: Callable[[], None] = None, + ) -> DurableChannel: + async def _on_closed(): + # fires once nats-py gives up reconnecting entirely - the closest + # analog to pika's channel-close callback, which today fires on + # any lost connection (this app doesn't implement reconnection). + on_channel_closed() + + self._nc = await nats.connect( + conf_get("nats_uri"), + closed_cb=_on_closed if on_channel_closed else None, + ) + self._js = self._nc.jetstream() + await _ensure_stream(self._js, queue) + + async def _cb(msg): + method = _Method(routing_key=msg.subject) + on_message(_JetStreamAckChannel(msg), method, _Properties(), msg.data) + + self._sub = await self._js.subscribe(queue, durable=_durable_name_for(queue), cb=_cb, manual_ack=True) + return _NatsDurableChannel(self._js, queue) + + async def close(self) -> None: + if self._nc is not None: + await self._nc.close() diff --git a/src/distbot/common/config/local_config.ini.spec b/src/distbot/common/config/local_config.ini.spec index 93f649e..a638ea0 100644 --- a/src/distbot/common/config/local_config.ini.spec +++ b/src/distbot/common/config/local_config.ini.spec @@ -14,3 +14,5 @@ hist_max_count = integer(default=5) hist_max_time = integer(default=10*60) amqp_uri = string(default="amqp://guest:guest@localhost:5672/%2F") +broker_backend = option("amqp", "nats", default="amqp") +nats_uri = string(default="nats://localhost:4222") diff --git a/src/distbot/common/message.py b/src/distbot/common/message.py index c8e22d7..5365c83 100644 --- a/src/distbot/common/message.py +++ b/src/distbot/common/message.py @@ -5,8 +5,7 @@ import logging import shlex from typing import List -import pika -from distbot.common.config import conf_get +from distbot.common.broker import get_sync_broker from slixmpp import Message, Presence from slixmpp.jid import JID @@ -45,19 +44,10 @@ def get_nick_from_message(message_obj: Message | Presence | dict) -> str: def process_message(routing_key, body: str): - connection = pika.BlockingConnection(pika.URLParameters(conf_get("amqp_uri"))) - channel = connection.channel() - channel.exchange_declare(exchange='classifier', exchange_type='topic') - - connection = pika.BlockingConnection( - pika.URLParameters(conf_get("amqp_uri")) - ) - channel = connection.channel() - logger.debug("Processing message body with routing key {}".format(routing_key)) - channel.basic_publish( - exchange='classifier', - routing_key=routing_key, - body=body - ) - connection.close() + broker = get_sync_broker() + broker.connect() + try: + broker.publish(routing_key, body) + finally: + broker.close() diff --git a/src/distbot/common/routing.py b/src/distbot/common/routing.py new file mode 100644 index 0000000..fbe6596 --- /dev/null +++ b/src/distbot/common/routing.py @@ -0,0 +1,15 @@ +# -*- coding: utf-8 -*- +import re + + +def matches(binding_key: str, subject: str) -> bool: + """Match a subject against an AMQP-topic-exchange-style binding key. + + '.' separates words, '*' matches exactly one word, '#' matches zero or + more words (including at the end, where a naive translation would + require at least one word). + """ + regex = binding_key.replace('.', r'\.').replace('*', '[^.]+').replace('#', '.*') + # fix .# leading to \..* + regex = re.sub(r"\\\.\.\*$", ".*", regex) + return re.fullmatch(regex, subject) is not None diff --git a/src/distbot/minijobber/run.py b/src/distbot/minijobber/run.py index f065ea0..3301a30 100644 --- a/src/distbot/minijobber/run.py +++ b/src/distbot/minijobber/run.py @@ -3,14 +3,12 @@ import logging import signal from time import sleep -import pika - from distbot.bot.worker import Worker from distbot.bot import worker as worker_mod -from distbot.common.config import conf_get from distbot.plugins import ( basic, fun, lookup, url, feeds, muc, translation, searx, queue_management, plugin_help, morse, meta, + vote, extended, bugtracker, bots, bofh, didyouknow, votepoll, youtube, @@ -42,6 +40,7 @@ PLUGIN_MODULES = { url: url.ALL, didyouknow: didyouknow.ALL, youtube: youtube.ALL, + vote: vote.ALL, # debug: debug.ALL votepoll: votepoll.ALL, } @@ -49,15 +48,6 @@ job_workers: list[Worker] = [] def initialize_workers(): - connection = pika.BlockingConnection( - pika.URLParameters(conf_get("amqp_uri")) - ) - channel = connection.channel() - channel.exchange_declare(exchange='topic_command', exchange_type='topic') - channel.exchange_declare(exchange='topic_parse', exchange_type='topic') - channel.queue_declare(queue=WORKER_QUEUE, durable=True) - channel.queue_declare(queue=ACTION_QUEUE, durable=True) - for classes in PLUGIN_MODULES.values(): for cls in classes: try: diff --git a/src/distbot/plugins/plugin_help.py b/src/distbot/plugins/plugin_help.py index cb9cb29..64ed45c 100644 --- a/src/distbot/plugins/plugin_help.py +++ b/src/distbot/plugins/plugin_help.py @@ -30,7 +30,7 @@ class Plugins(Worker): def init_queue(self): super().init_queue() - self.channel.basic_consume(queue='plugin_registry', on_message_callback=self.callback_plugins) + self.broker.register_queue_consumer('plugin_registry', self.callback_plugins, max_age_seconds=120) def callback_plugins(self, ch, method, properties, body): body = json.loads(body.decode("utf-8")) diff --git a/tests/test_integration/bot_test_utils.py b/tests/test_integration/bot_test_utils.py index 64fca83..1af3d48 100644 --- a/tests/test_integration/bot_test_utils.py +++ b/tests/test_integration/bot_test_utils.py @@ -1,5 +1,4 @@ # -*- coding: utf-8 -*- -import re from dataclasses import dataclass from unittest import mock @@ -10,6 +9,7 @@ from slixmpp import Message, JID from distbot.bot.bot import Bot from distbot.bot.worker import Worker +from distbot.common.routing import matches @dataclass @@ -51,21 +51,14 @@ class LocalMessageBroker: params = get_amqp_callback_params(msg, sender, bot_nick) for recipient in self.recipients: for binding_key in recipient.binding_keys: - if self.matches(binding_key, params.method.routing_key): + if matches(binding_key, params.method.routing_key): recipient.callback(**params.__dict__) break - @staticmethod - def matches(binding_key, message_routing_key): - regex = binding_key.replace('.', '\\.').replace('*', '[^.]+').replace('#', '.*') - # fix .# leading to \..* - regex = re.sub(r"\\.\.\*$", ".*", regex) - return re.fullmatch(regex, message_routing_key) is not None - @pytest.mark.parametrize("bind,msg_routing_key", [ ("nick.dice.#", "nick.dice"), ("nick.dice.#", "nick.dice.2"), ]) def test_local_message_broker(bind, msg_routing_key): - assert LocalMessageBroker.matches(bind, msg_routing_key) \ No newline at end of file + assert matches(bind, msg_routing_key) \ No newline at end of file diff --git a/tests/test_unit/test_action_worker.py b/tests/test_unit/test_action_worker.py new file mode 100644 index 0000000..ea4200c --- /dev/null +++ b/tests/test_unit/test_action_worker.py @@ -0,0 +1,206 @@ +# -*- coding: utf-8 -*- +import asyncio +import json +import time +from unittest.mock import Mock + +import pytest + +import distbot.bot.action_worker as aw_module +from distbot.bot.action_worker import ActionWorker +from distbot.common.action import Action + + +def _in(seconds=60, **kwargs): + return Action(time=time.time() + seconds, **kwargs) + + +@pytest.fixture(autouse=True) +def mock_conf(monkeypatch): + monkeypatch.setattr(aw_module, 'conf_get', Mock(return_value='0')) + monkeypatch.setattr(aw_module, 'conf_set', Mock()) + + +@pytest.fixture(autouse=True) +def no_sleep(monkeypatch): + monkeypatch.setattr(aw_module.time, 'sleep', Mock()) + + +@pytest.fixture() +def mock_process_message(monkeypatch): + m = Mock() + monkeypatch.setattr(aw_module, 'process_message', m) + return m + + +@pytest.fixture() +def aw(monkeypatch): + loop = asyncio.new_event_loop() + monkeypatch.setattr(asyncio, 'get_event_loop', lambda: loop) + + def _create_task(coro): + coro.close() # close coroutine without running it to avoid warning + return Mock() + + original_create_task = loop.create_task + loop.create_task = _create_task + worker = ActionWorker(bot=Mock(), queue_name='test_actions') + loop.create_task = original_create_task # run_until_complete() needs the real one + worker.event_list = [] # give each test its own list (class attr is shared) + worker._channel = Mock() # intercept _publish_action calls + yield worker + loop.close() + + +# --- schedule_action: timer fires on channel --- + +def test_schedule_fires_publish_on_channel(aw): + aw._channel = Mock() + action = _in(seconds=0, mutex='fire') + aw.schedule_action(action) + aw.loop.run_until_complete(asyncio.sleep(0)) + aw._channel.basic_publish.assert_called_once_with( + exchange='', + routing_key='test_actions', + body=action.serialize() + ) + + +# --- schedule_action --- + +def test_schedule_adds_handle(aw): + aw.schedule_action(_in(mutex='t1')) + assert len(aw.event_list) == 1 + + +def test_schedule_blocks_duplicate_mutex(aw): + aw.schedule_action(_in(mutex='t1')) + with pytest.raises(RuntimeError): + aw.schedule_action(_in(mutex='t1')) + + +def test_schedule_allows_different_mutexes(aw): + aw.schedule_action(_in(mutex='t1')) + aw.schedule_action(_in(mutex='t2')) + assert len(aw.event_list) == 2 + + +def test_schedule_without_mutex_always_allowed(aw): + aw.schedule_action(_in()) + aw.schedule_action(_in()) + assert len(aw.event_list) == 2 + + +# --- unschedule_action --- + +def test_unschedule_cancels_handle(aw): + aw.schedule_action(_in(mutex='stopper')) + handle = aw.event_list[0] + + aw.unschedule_action(Action(mutex='stopper')) + + assert handle.cancelled() + + +def test_unschedule_unknown_mutex_is_noop(aw): + aw.schedule_action(_in(mutex='exists')) + aw.unschedule_action(Action(mutex='ghost')) + assert not aw.event_list[0].cancelled() + + +def test_unschedule_allows_reschedule_after_cancel(aw): + aw.schedule_action(_in(mutex='reuse')) + aw.unschedule_action(Action(mutex='reuse')) + aw.schedule_action(_in(mutex='reuse')) # should not raise + + +# --- run_action: event scheduling --- + +def test_run_action_schedules_nested_event(aw): + action = Action(msg='set timer', recipient='room@conf', sender='user', event=_in(mutex='ev1')) + aw.run_action(action) + assert len(aw.event_list) == 1 + + +def test_run_action_event_inherits_recipient_from_parent(aw): + event = _in(mutex='ev1') + action = Action(msg='hi', recipient='room@conf', sender='user', event=event) + aw.run_action(action) + scheduled = aw.event_list[0]._args[0] + assert scheduled.recipient == 'room@conf' + + +def test_run_action_event_inherits_sender_from_parent(aw): + event = _in(mutex='ev1') + action = Action(msg='hi', recipient='room@conf', sender='user@room', event=event) + aw.run_action(action) + scheduled = aw.event_list[0]._args[0] + assert scheduled.sender == 'user@room' + + +def test_run_action_event_keeps_explicit_recipient(aw): + event = _in(mutex='ev1', recipient='explicit@room') + action = Action(recipient='parent@room', sender='user', event=event) + aw.run_action(action) + scheduled = aw.event_list[0]._args[0] + assert scheduled.recipient == 'explicit@room' + + +# --- run_action: stop_event --- + +def test_run_action_stop_event_cancels_timer(aw): + aw.schedule_action(_in(mutex='cancel_me')) + handle = aw.event_list[0] + + action = Action(event=Action(stop_event=True, mutex='cancel_me')) + aw.run_action(action) + + assert handle.cancelled() + + +# --- run_action: mutex conflict warning --- + +def test_run_action_mutex_conflict_prepends_warning(aw): + aw.schedule_action(_in(mutex='conflict')) + action = Action(msg='original', event=_in(mutex='conflict')) + aw.run_action(action) + assert 'Warning' in action.msg + assert 'original' in action.msg + + +def test_run_action_no_warning_without_msg(aw): + aw.schedule_action(_in(mutex='conflict')) + action = Action(event=_in(mutex='conflict')) + aw.run_action(action) # should not raise, msg stays None + assert action.msg is None + + +# --- run_action: command replay --- + +def test_run_action_command_replays_to_broker(aw, mock_process_message): + action = Action(command='nick.dice.5', recipient='room@conf') + aw.run_action(action) + mock_process_message.assert_called_once() + call_kwargs = mock_process_message.call_args.kwargs + assert call_kwargs['routing_key'] == b'nick.dice.5' + + +def test_run_action_command_body_strips_nick_prefix(aw, mock_process_message): + action = Action(command='nick.dice.5', recipient='room@conf') + aw.run_action(action) + body = json.loads(mock_process_message.call_args.kwargs['body']) + assert body['body'] == 'dice 5' + + +# --- run_action: message delivery --- + +def test_run_action_sends_msg_via_bot(aw): + action = Action(msg='hello', recipient='room@conf') + aw.run_action(action) + aw._bot.echo.assert_called_once_with('hello', 'room@conf') + + +def test_run_action_sends_priv_msg_to_sender(aw): + action = Action(priv_msg='whisper', sender='user@room') + aw.run_action(action) + aw._bot.echo.assert_called_once_with('whisper', 'user@room') diff --git a/tests/test_unit/test_worker_callback.py b/tests/test_unit/test_worker_callback.py new file mode 100644 index 0000000..a4f743b --- /dev/null +++ b/tests/test_unit/test_worker_callback.py @@ -0,0 +1,115 @@ +# -*- coding: utf-8 -*- +import json +from unittest.mock import Mock + +import pytest + +import distbot.bot.worker as worker_module +from distbot.bot.worker import Worker +from distbot.common.action import Action + + +class _TestWorker(Worker): + binding_keys = ['nick.test.#'] + usage = "testbot: test" + returned_action = None + + def parse_body(self, msg): + return self.returned_action + + +@pytest.fixture() +def w(): + worker = _TestWorker(actionqueue='test_actions') + worker.channel = Mock() + return worker + + +@pytest.fixture() +def send_action(monkeypatch): + mock_action = Mock() + monkeypatch.setattr(worker_module, 'send_action', mock_action) + return mock_action + + +def _make_pika_args(routing_key='nick.test.foo', body_from='user@room/User', body_to='room@conf', body_text='test foo'): + ch = Mock() + method = Mock() + method.routing_key = routing_key + method.delivery_tag = 42 + properties = Mock() + body = json.dumps({'from': body_from, 'to': body_to, 'body': body_text}).encode() + return ch, method, properties, body + + +def _published_action(send_action): + args = send_action.call_args.args + assert args[0] == 'test_actions' + return args[1] + + +def test_callback_publishes_action(w, send_action): + w.returned_action = Action(msg='hello') + ch, method, properties, body = _make_pika_args() + + w.callback(ch, method, properties, body) + + send_action.assert_called_once() + assert _published_action(send_action).msg == 'hello' + + +def test_callback_sets_sender_and_recipient(w, send_action): + w.returned_action = Action(msg='hi') + ch, method, properties, body = _make_pika_args(body_from='alice@room/Alice', body_to='room@conf') + + w.callback(ch, method, properties, body) + + action = _published_action(send_action) + assert action.sender == 'alice@room/Alice' + assert action.recipient == 'room@conf' + + +def test_callback_preserves_explicit_recipient(w, send_action): + w.returned_action = Action(msg='hi', recipient='custom@room') + ch, method, properties, body = _make_pika_args(body_to='other@room') + + w.callback(ch, method, properties, body) + + assert _published_action(send_action).recipient == 'custom@room' + + +def test_callback_no_publish_when_no_action(w, send_action): + w.returned_action = None + ch, method, properties, body = _make_pika_args() + + w.callback(ch, method, properties, body) + + send_action.assert_not_called() + + +def test_callback_acks_on_action(w, send_action): + w.returned_action = Action(msg='hi') + ch, method, properties, body = _make_pika_args() + + w.callback(ch, method, properties, body) + + ch.basic_ack.assert_called_once_with(delivery_tag=42) + + +def test_callback_acks_when_no_action(w): + w.returned_action = None + ch, method, properties, body = _make_pika_args() + + w.callback(ch, method, properties, body) + + ch.basic_ack.assert_called_once_with(delivery_tag=42) + + +def test_callback_acks_even_if_parse_body_raises(w): + w.parse_body = Mock(side_effect=RuntimeError("plugin exploded")) + ch, method, properties, body = _make_pika_args() + + w.callback(ch, method, properties, body) + + w.channel.basic_publish.assert_not_called() + ch.basic_ack.assert_called_once_with(delivery_tag=42)