comet-ml--opik/sdks/python/tests/unit/message_processing/test_message_queue.py

import threading
import time

import pytest

from opik.message_processing import message_queue


def test_message_queue_accept_put_without_discarding___max_size_handling():
    max_size = 10
    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue(max_size)

    queue.put("the oldest message")
    assert queue.accept_put_without_discarding() is True

    for i in range(max_size - 1):
        assert queue.accept_put_without_discarding() is True, f"Failed at {i}"
        queue.put(f"line: {i}")

    assert queue.accept_put_without_discarding() is False

    # put one more message to have the oldest discarded
    queue.put("the newest message")

    # check that the oldest message ("the oldest message") was discarded
    assert queue.get(0.0001) == "line: 0"


def test_message_queue_accept_put_without_discarding___max_size_none():
    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue(None)
    for i in range(100):
        queue.put(f"line: {i}")

    assert queue.accept_put_without_discarding() is True


def test_message_queue_task_done__fresh_queue__reports_all_tasks_done():
    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()
    assert queue.all_tasks_done() is True


def test_message_queue_task_done__put_then_task_done__quiescence_restored():
    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()

    queue.put("msg-1")
    assert queue.all_tasks_done() is False

    queue.get(0.0001)
    # Popping does not mark the task done — it's in flight until task_done().
    assert queue.all_tasks_done() is False

    queue.task_done()
    assert queue.all_tasks_done() is True


def test_message_queue_task_done__multiple_puts__decremented_one_by_one():
    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()

    for i in range(5):
        queue.put(f"msg-{i}")

    assert queue.all_tasks_done() is False

    for _ in range(4):
        queue.get(0.0001)
        queue.task_done()
        assert queue.all_tasks_done() is False

    queue.get(0.0001)
    queue.task_done()
    assert queue.all_tasks_done() is True


def test_message_queue_task_done__put_back_does_not_increment_counter():
    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()

    queue.put("msg")
    msg = queue.get(0.0001)
    # Re-enqueue (delivery-time deferral / rate-limit retry): the task is
    # still in flight, so the counter must not change.
    queue.put_back(msg)
    assert queue.all_tasks_done() is False

    queue.get(0.0001)
    queue.task_done()
    assert queue.all_tasks_done() is True


def test_message_queue_task_done__called_too_many_times__raises():
    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()

    queue.put("msg")
    queue.get(0.0001)
    queue.task_done()

    with pytest.raises(ValueError):
        queue.task_done()


def test_message_queue_task_done__maxlen_eviction__does_not_increment_counter():
    # When a bounded queue is full, appendleft silently evicts the oldest
    # item. The displaced message's task carries over to the newcomer, so
    # the unfinished-task count must not grow.
    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue(max_length=2)

    queue.put("a")
    queue.put("b")
    # Now full — this put evicts the oldest end.
    queue.put("c")

    # Only two messages are reachable; only two task_done() calls should
    # bring the queue back to quiescence.
    queue.get(0.0001)
    queue.task_done()
    queue.get(0.0001)
    queue.task_done()
    assert queue.all_tasks_done() is True


def test_message_queue_put_back__bounded_queue_full__evicted_task_is_released():
    # Regression: on a bounded queue, `put_back` (used by the rate-limit
    # retry path) silently evicts an item from the opposite end. That
    # evicted item was already counted in `_unfinished_tasks` at its
    # original `put()` time and will never see `task_done()` — so the
    # counter must drop by 1 to avoid stalling `join()`.
    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue(max_length=2)

    queue.put("a")  # unfinished=1
    queue.put("b")  # unfinished=2; deque [b, a]

    # Simulate a consumer popping a message and being rate-limited mid-process.
    msg = queue.get(0.0001)  # deque [b]; unfinished=2 (msg still in flight)
    assert msg == "a"

    queue.put("c")  # unfinished=3; deque [c, b]

    # Re-enqueue the popped message: the deque is full, so `append` evicts
    # "c" from the left. "c" will never be processed, so its task slot must
    # be released.
    queue.put_back(msg)

    # Two reachable messages — two task_done() calls must restore quiescence.
    queue.get(0.0001)
    queue.task_done()
    queue.get(0.0001)
    queue.task_done()
    assert queue.all_tasks_done() is True
    assert queue.join(timeout=0.05) is True


def test_message_queue_clear__resets_counter_and_unblocks_join():
    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()

    for i in range(3):
        queue.put(f"msg-{i}")

    assert queue.all_tasks_done() is False

    queue.clear()

    # clear() drops messages that will never see task_done(); the counter
    # must be reset, otherwise join()/all_tasks_done() would stall forever.
    assert queue.all_tasks_done() is True
    assert queue.join(timeout=0.05) is True


def test_message_queue_join__no_pending_tasks__returns_true_immediately():
    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()
    start = time.monotonic()
    assert queue.join(timeout=1.0) is True
    assert time.monotonic() - start < 0.1


def test_message_queue_join__pending_tasks_with_timeout__returns_false():
    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()
    queue.put("msg")

    start = time.monotonic()
    assert queue.join(timeout=0.1) is False
    elapsed = time.monotonic() - start
    assert 0.1 <= elapsed < 0.5


def test_message_queue_join__task_done_in_other_thread__wakes_waiter():
    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()
    queue.put("msg")
    queue.get(0.0001)

    def _consumer() -> None:
        time.sleep(0.05)
        queue.task_done()

    threading.Thread(target=_consumer, daemon=True).start()

    start = time.monotonic()
    assert queue.join(timeout=2.0) is True
    elapsed = time.monotonic() - start
    # Should wake up shortly after the 0.05s sleep, not wait the full timeout.
    assert elapsed < 1.0


@pytest.mark.parametrize(
    "maximal_queue_size,batch_factor,expected",
    [
        (100, 10, 10),
        (100, 0, 100),
    ],
)
def test_calculate_max_queue_size(maximal_queue_size, batch_factor, expected):
    assert (
        message_queue.calculate_max_queue_size(maximal_queue_size, batch_factor)
        == expected
    )