iTi-Flask/tests/test_mq.py

import asyncio
import builtins
import sys
from pathlib import Path

import pytest

from iti import create_app
from iti.config import BaseConfig
from iti.mq import MQConfigError, init_mq, mq_client, mq_consumer
from iti.mq.backend import KafkaBackend
from iti.mq.offset_store import MemoryMQOffsetStore, SQLiteMQOffsetStore
from iti.mq.registry import MQRegistry as RegistryClass


class FakeProducer:
    def __init__(self, config):
        self.config = config
        self.produced = []
        self.polled = []
        self.flushed = []

    def poll(self, timeout):
        self.polled.append(timeout)

    def produce(self, topic, *, value=None, key=None, headers=None):
        self.produced.append(
            {"topic": topic, "value": value, "key": key, "headers": headers}
        )

    def flush(self, timeout=None):
        self.flushed.append(timeout)


class FakeConsumer:
    def __init__(self, messages=None):
        self.messages = list(messages or [])
        self.subscribed = []
        self.assigned = []
        self.committed = []
        self.sought = []
        self.closed = False
        self.config = None
        self.raise_offset_out_of_range_once = False

    def subscribe(self, topics):
        self.subscribed.append(topics)

    def assign(self, partitions):
        self.assigned.append(
            [
                (
                    call_or_value(partition, "topic"),
                    call_or_value(partition, "partition"),
                    call_or_value(partition, "offset"),
                )
                for partition in partitions
            ]
        )

    def poll(self, timeout):
        if self.raise_offset_out_of_range_once:
            self.raise_offset_out_of_range_once = False
            return FakeErrorMessage(topic="demo.topic", partition=0)
        if self.messages:
            return self.messages.pop(0)
        return None

    def commit(self, message, asynchronous=False):
        self.committed.append((message, asynchronous))

    def seek(self, position):
        if all(hasattr(position, name) for name in ("topic", "partition", "offset")):
            self.sought.append(
                (
                    call_or_value(position, "topic"),
                    call_or_value(position, "partition"),
                    call_or_value(position, "offset"),
                )
            )
            return
        self.sought.append(position)

    def close(self):
        self.closed = True

    def list_topics(self, timeout=10):
        return FakeMetadata(
            {
                "demo.topic": FakeTopicMetadata([0, 1]),
                "demo.other": FakeTopicMetadata([2]),
            }
        )

    def get_watermark_offsets(self, partition, timeout=10, cached=False):
        return (5, 11)


class FakeKafkaError:
    _OFFSET_OUT_OF_RANGE = 1

    def __init__(self, code=1, message="offset out of range"):
        self._code = code
        self._message = message

    def code(self):
        return self._code

    def __str__(self):
        return self._message


class FakeErrorMessage:
    def __init__(self, *, topic="demo.topic", partition=0):
        self._topic = topic
        self._partition = partition
        self._error = FakeKafkaError()

    def error(self):
        return self._error

    def topic(self):
        return self._topic

    def partition(self):
        return self._partition


class FakeTopicPartition:
    def __init__(self, topic, partition, offset=None):
        self.topic = topic
        self.partition = partition
        self.offset = offset


class FakeTopicMetadata:
    def __init__(self, partitions, error=None):
        self.partitions = {partition: object() for partition in partitions}
        self.error = error


class FakeMetadata:
    def __init__(self, topics):
        self.topics = topics


class FakeMessage:
    def __init__(
        self,
        *,
        topic="demo.topic",
        partition=0,
        offset=1,
        key=b"k1",
        value=b'{"ok":true}',
        headers=None,
    ):
        self._topic = topic
        self._partition = partition
        self._offset = offset
        self._key = key
        self._value = value
        self._headers = headers or [("source", b"test")]

    def topic(self):
        return self._topic

    def partition(self):
        return self._partition

    def offset(self):
        return self._offset

    def key(self):
        return self._key

    def value(self):
        return self._value

    def headers(self):
        return self._headers

    def timestamp(self):
        return (0, 0)

    def error(self):
        return None


def test_mq_registry_registers_decorator_and_explicit_consumer():
    registry = RegistryClass()
    registry.register_consumer(
        name="explicit",
        topics=["demo.explicit"],
        group_id="g1",
        handler=lambda message: None,
    )

    assert registry.consumers["explicit"].topics == ("demo.explicit",)

    before = set(mq_consumer_registry_names())

    try:
        @mq_consumer("demo.decorated", name="decorated-test", group_id="g1")
        def decorated(message):
            return None

        assert "decorated-test" in mq_consumer_registry_names() - before
        assert decorated.__name__ == "decorated"
    finally:
        from iti.mq import mq_registry

        mq_registry.consumers.pop("decorated-test", None)


def test_mq_registry_rejects_duplicate_names():
    registry = RegistryClass()
    registry.register_producer(name="events", topic="demo.events")
    registry.register_consumer(
        name="consumer",
        topics="demo.events",
        handler=lambda message: None,
    )

    with pytest.raises(ValueError, match="producer already registered"):
        registry.register_producer(name="events", topic="other")
    with pytest.raises(ValueError, match="consumer already registered"):
        registry.register_consumer(
            name="consumer",
            topics="other",
            handler=lambda message: None,
        )


def test_mq_enabled_false_does_not_import_or_configure_kafka(monkeypatch):
    sys.modules.pop("confluent_kafka", None)

    app = create_app(
        config_mapping=BaseConfig(
            database_url="sqlite+pysqlite:///:memory:",
            testing=True,
            mq_enabled=False,
            exchange_enabled=False,
        )
    )

    assert not hasattr(app.state, "iti_mq_client")
    assert "confluent_kafka" not in sys.modules


def test_mq_enabled_true_without_dependency_raises_install_hint(monkeypatch):
    sys.modules.pop("confluent_kafka", None)
    real_import = builtins.__import__

    def missing_confluent_kafka(name, *args, **kwargs):
        if name == "confluent_kafka":
            raise ImportError("missing")
        return real_import(name, *args, **kwargs)

    monkeypatch.setattr(builtins, "__import__", missing_confluent_kafka)

    with pytest.raises(MQConfigError, match=r"iti-flask\[mq-kafka\]"):
        create_app(
            config_mapping=BaseConfig(
                database_url="sqlite+pysqlite:///:memory:",
                testing=True,
                mq_enabled=True,
                mq={"backend": "kafka", "bootstrap_servers": "127.0.0.1:9092"},
                exchange_enabled=False,
            )
        )


def test_mq_client_sends_json_bytes_and_registered_sender():
    producer = FakeProducer({"bootstrap.servers": "localhost:9092"})
    registry = RegistryClass()
    registry.register_producer(name="events", topic="demo.events")
    app = create_app(
        config_mapping=BaseConfig(
            database_url="sqlite+pysqlite:///:memory:",
            testing=True,
            mq_enabled=False,
            exchange_enabled=False,
        )
    )
    init_mq(
        app,
        {"backend": "kafka", "bootstrap_servers": "localhost:9092"},
        registry=registry,
        producer_factory=lambda config: producer,
    )

    client = mq_client(app)
    client.send_json("demo.raw", {"id": "1"}, key="k1", headers={"h": "v"})
    client.send("demo.bytes", value=b"raw", key=b"k2")
    client.sender("events").send_json({"id": "2"})
    client.flush(2)

    assert producer.produced[0] == {
        "topic": "demo.raw",
        "value": b'{"id":"1"}',
        "key": b"k1",
        "headers": [("h", b"v")],
    }
    assert producer.produced[1]["value"] == b"raw"
    assert producer.produced[1]["key"] == b"k2"
    assert producer.produced[2]["topic"] == "demo.events"
    assert producer.flushed == [2]


def test_mq_client_supports_inline_flush():
    producer = FakeProducer({"bootstrap.servers": "localhost:9092"})
    app = create_app(
        config_mapping=BaseConfig(
            database_url="sqlite+pysqlite:///:memory:",
            testing=True,
            mq_enabled=False,
            exchange_enabled=False,
        )
    )
    init_mq(
        app,
        {"backend": "kafka", "bootstrap_servers": "localhost:9092"},
        registry=RegistryClass(),
        producer_factory=lambda config: producer,
    )

    mq_client(app).send_json("demo.flush", {"id": "1"}, flush=True)

    assert producer.flushed == [None]


def test_runner_commits_after_successful_sync_handler():
    handled = []
    message = FakeMessage()
    fake_consumer = FakeConsumer([message])
    app = make_mq_app(
        registry_with_consumer(lambda item: handled.append(item.value)),
        fake_consumer,
    )
    runner = app.state.iti_mq_runner

    runner.start()
    wait_until(lambda: bool(fake_consumer.committed))
    runner.stop()

    assert handled == [{"ok": True}]
    assert fake_consumer.committed == [(message, False)]
    assert fake_consumer.sought == []
    assert fake_consumer.closed is True


def test_runner_executes_async_handler():
    handled = []

    async def handler(message):
        await asyncio.sleep(0)
        handled.append(message.key)

    fake_consumer = FakeConsumer([FakeMessage()])
    app = make_mq_app(registry_with_consumer(handler), fake_consumer)
    runner = app.state.iti_mq_runner

    runner.start()
    wait_until(lambda: bool(fake_consumer.committed))
    runner.stop()

    assert handled == ["k1"]


def test_runner_does_not_commit_and_seeks_after_handler_failure():
    message = FakeMessage()
    fake_consumer = FakeConsumer([message])

    def handler(_message):
        raise RuntimeError("boom")

    app = make_mq_app(registry_with_consumer(handler), fake_consumer)
    runner = app.state.iti_mq_runner

    runner.start()
    wait_until(lambda: bool(fake_consumer.sought))
    runner.stop()

    assert fake_consumer.committed == []
    assert fake_consumer.sought == [("demo.topic", 0, 1)]


def test_runner_raises_when_group_id_missing():
    registry = RegistryClass()
    registry.register_consumer(
        name="demo",
        topics="demo.topic",
        handler=lambda message: None,
    )
    app = make_mq_app(registry, FakeConsumer())

    with pytest.raises(MQConfigError, match="missing group_id"):
        app.state.iti_mq_runner.start()


def test_kafka_backend_forces_manual_commit():
    backend = KafkaBackend(
        {
            "bootstrap_servers": "localhost:9092",
            "group_id": "global",
            "auto_offset_reset": "latest",
            "consumer": {"enable.auto.commit": True},
        },
        producer_factory=lambda config: FakeProducer(config),
        consumer_factory=lambda config: FakeConsumer(),
    )

    config = backend.consumer_config("g1")

    assert config["bootstrap.servers"] == "localhost:9092"
    assert config["group.id"] == "g1"
    assert config["auto.offset.reset"] == "latest"
    assert config["enable.auto.commit"] is False


def test_runner_assign_mode_uses_topic_metadata_and_offset_store(monkeypatch, tmp_path):
    monkeypatch_confluent(monkeypatch)
    handled = []
    fake_consumer = FakeConsumer([FakeMessage(topic="demo.topic", partition=1, offset=7)])
    registry = RegistryClass()
    registry.register_consumer(
        name="assign-demo",
        topics="demo.topic",
        mode="assign",
        handler=lambda message: handled.append((message.partition, message.offset)),
    )

    app = create_app(
        config_mapping=BaseConfig(
            database_url="sqlite+pysqlite:///:memory:",
            testing=True,
            exchange_enabled=False,
        )
    )
    init_mq(
        app,
        {
            "backend": "kafka",
            "bootstrap_servers": "localhost:9092",
            "consumer_mode": "assign",
            "offset_store": {"type": "memory"},
            "failure_backoff_seconds": 0.01,
            "poll_timeout_seconds": 0.01,
        },
        registry=registry,
        producer_factory=lambda config: FakeProducer(config),
        consumer_factory=lambda config: fake_consumer,
    )

    runner = app.state.iti_mq_runner
    runner.start()
    wait_until(lambda: handled == [(1, 7)])
    store = runner._workers[0].offset_store
    runner.stop()

    assert fake_consumer.assigned == [[("demo.topic", 0, -2), ("demo.topic", 1, -2)]]
    assert isinstance(store, MemoryMQOffsetStore)
    assert store.get("assign-demo", "demo.topic", 1) == 8


def test_runner_assign_mode_uses_global_consumer_mode(monkeypatch):
    monkeypatch_confluent(monkeypatch)
    fake_consumer = FakeConsumer()
    registry = RegistryClass()
    registry.register_consumer(
        name="assign-global",
        topics="demo.topic",
        handler=lambda message: None,
    )
    app = make_mq_app(
        registry,
        fake_consumer,
        config_overrides={
            "consumer_mode": "assign",
            "offset_store": {"type": "memory"},
        },
    )

    app.state.iti_mq_runner.start()
    app.state.iti_mq_runner.stop()

    assert fake_consumer.assigned == [[("demo.topic", 0, -2), ("demo.topic", 1, -2)]]
    assert fake_consumer.subscribed == []


def test_runner_assign_mode_rejects_invalid_offset_store(monkeypatch):
    monkeypatch_confluent(monkeypatch)
    fake_consumer = FakeConsumer()
    registry = RegistryClass()
    registry.register_consumer(
        name="assign-invalid-store",
        topics="demo.topic",
        mode="assign",
        handler=lambda message: None,
    )
    app = make_mq_app(
        registry,
        fake_consumer,
        config_overrides={
            "consumer_mode": "assign",
            "offset_store": {"type": "unsupported"},
        },
    )

    with pytest.raises(MQConfigError, match="unsupported mq offset store type"):
        app.state.iti_mq_runner.start()


def test_runner_assign_mode_recovers_from_offset_out_of_range(monkeypatch):
    monkeypatch_confluent(monkeypatch)
    fake_consumer = FakeConsumer([FakeMessage(topic="demo.topic", partition=0, offset=5)])
    fake_consumer.raise_offset_out_of_range_once = True
    registry = RegistryClass()
    registry.register_consumer(
        name="assign-recover",
        topics="demo.topic",
        mode="assign",
        handler=lambda message: None,
    )
    app = make_mq_app(
        registry,
        fake_consumer,
        config_overrides={
            "consumer_mode": "assign",
            "offset_store": {"type": "memory"},
            "auto_offset_reset": "latest",
        },
    )

    runner = app.state.iti_mq_runner
    runner.start()
    wait_until(lambda: ("demo.topic", 0, 11) in fake_consumer.sought)
    runner.stop()

    assert ("demo.topic", 0, 11) in fake_consumer.sought


def test_sqlite_offset_store_round_trip(tmp_path):
    path = tmp_path / "mq-offsets.sqlite"
    store = SQLiteMQOffsetStore(path)
    try:
        store.set("demo", "topic", 1, 9)
        assert store.get("demo", "topic", 1) == 9
    finally:
        store.close()
    assert Path(path).exists()


def mq_consumer_registry_names():
    from iti.mq import mq_registry

    return set(mq_registry.consumers.keys())


def registry_with_consumer(handler):
    registry = RegistryClass()
    registry.register_consumer(
        name="demo",
        topics="demo.topic",
        group_id="g1",
        handler=handler,
    )
    return registry


def make_mq_app(registry, fake_consumer, config_overrides=None):
    app = create_app(
        config_mapping=BaseConfig(
            database_url="sqlite+pysqlite:///:memory:",
            testing=True,
            mq_enabled=False,
            exchange_enabled=False,
        )
    )

    def consumer_factory(config):
        fake_consumer.config = config
        return fake_consumer

    config = {
        "backend": "kafka",
        "bootstrap_servers": "localhost:9092",
        "failure_backoff_seconds": 0.01,
        "poll_timeout_seconds": 0.01,
    }
    if config_overrides:
        config.update(config_overrides)

    init_mq(
        app,
        config,
        registry=registry,
        producer_factory=lambda config: FakeProducer(config),
        consumer_factory=consumer_factory,
    )
    return app


def wait_until(predicate, timeout=1.0):
    loop = asyncio.new_event_loop()
    try:
        deadline = loop.time() + timeout
        while loop.time() < deadline:
            if predicate():
                return
            loop.run_until_complete(asyncio.sleep(0.01))
    finally:
        loop.close()
    raise AssertionError("condition not reached")


def call_or_value(value, name):
    attr = getattr(value, name)
    return attr() if callable(attr) else attr


def monkeypatch_confluent(monkeypatch):
    import iti.mq.runner as runner_module

    monkeypatch.setattr(runner_module, "_auto_offset", lambda runner: -1 if runner.backend.config.get("auto_offset_reset") == "latest" else -2)
    monkeypatch.setitem(sys.modules, "confluent_kafka", type("FakeConfluent", (), {
        "TopicPartition": FakeTopicPartition,
        "KafkaError": FakeKafkaError,
        "OFFSET_BEGINNING": -2,
        "OFFSET_END": -1,
    }))