fed8b2eed7
Backend release / release (push) Waiting to run
Bandit Security Scan / bandit_scan (push) Waiting to run
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Waiting to run
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Waiting to run
Build and push multi-arch DocsGPT Docker image / manifest (push) Blocked by required conditions
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Waiting to run
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Waiting to run
Build and push DocsGPT FE Docker image for development / manifest (push) Blocked by required conditions
Python linting / ruff (push) Waiting to run
Run python tests with pytest / Run tests and count coverage (3.12) (push) Waiting to run
React Widget Build / build (push) Waiting to run
847 lines
38 KiB
Python
847 lines
38 KiB
Python
"""Unit tests for ``application/streaming/message_journal.py``.
|
||
|
||
The journal hook is best-effort by contract — its failure modes are
|
||
the most important thing to lock down so a streaming hiccup never
|
||
crashes ``complete_stream``.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
import pytest
|
||
|
||
from application.streaming.message_journal import (
|
||
BatchedJournalWriter,
|
||
record_event,
|
||
)
|
||
|
||
|
||
def _integrity_error(sqlstate: str):
|
||
"""Build an ``IntegrityError`` whose ``.orig`` reports ``sqlstate``.
|
||
|
||
Mirrors how psycopg3 surfaces a Postgres error code on the DBAPI
|
||
exception that SQLAlchemy wraps: ``23503`` = foreign-key violation
|
||
(the ``conversation_messages`` parent row is missing), ``23505`` =
|
||
unique/PK collision (a duplicate ``sequence_no``).
|
||
"""
|
||
from sqlalchemy.exc import IntegrityError
|
||
|
||
orig = Exception("pg error")
|
||
orig.sqlstate = sqlstate
|
||
return IntegrityError("stmt", {}, orig)
|
||
|
||
|
||
@pytest.mark.unit
|
||
class TestRecordEvent:
|
||
def test_invalid_args_return_false(self):
|
||
assert record_event("", 0, "answer") is False
|
||
assert record_event("msg-1", 0, "") is False
|
||
|
||
def test_happy_path_writes_and_publishes(self):
|
||
with patch(
|
||
"application.streaming.message_journal.db_session"
|
||
) as mock_session, patch(
|
||
"application.streaming.message_journal.MessageEventsRepository"
|
||
) as mock_repo_cls, patch(
|
||
"application.streaming.message_journal.Topic"
|
||
) as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo = MagicMock()
|
||
mock_repo_cls.return_value = mock_repo
|
||
mock_topic = MagicMock()
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
result = record_event(
|
||
"msg-1", 5, "answer", {"type": "answer", "answer": "ok"}
|
||
)
|
||
|
||
assert result is True
|
||
mock_repo.record.assert_called_once_with(
|
||
"msg-1", 5, "answer", {"type": "answer", "answer": "ok"}
|
||
)
|
||
mock_topic_cls.assert_called_once_with("channel:msg-1")
|
||
mock_topic.publish.assert_called_once()
|
||
wire = mock_topic.publish.call_args[0][0]
|
||
envelope = json.loads(wire)
|
||
assert envelope["sequence_no"] == 5
|
||
assert envelope["event_type"] == "answer"
|
||
assert envelope["payload"] == {"type": "answer", "answer": "ok"}
|
||
|
||
def test_publish_attempted_even_when_journal_fails(self):
|
||
"""A DB hiccup must not stop the live tail — currently-attached
|
||
subscribers should still receive the live event so their UI is
|
||
live even if a future reconnect's snapshot is missing this row.
|
||
"""
|
||
with patch(
|
||
"application.streaming.message_journal.db_session"
|
||
) as mock_session, patch(
|
||
"application.streaming.message_journal.MessageEventsRepository"
|
||
), patch(
|
||
"application.streaming.message_journal.Topic"
|
||
) as mock_topic_cls:
|
||
mock_session.return_value.__enter__.side_effect = RuntimeError("pg down")
|
||
mock_topic = MagicMock()
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
result = record_event(
|
||
"msg-1", 1, "answer", {"type": "answer", "answer": "x"}
|
||
)
|
||
|
||
assert result is False
|
||
mock_topic.publish.assert_called_once()
|
||
|
||
def test_publish_failure_does_not_raise(self):
|
||
with patch(
|
||
"application.streaming.message_journal.db_session"
|
||
) as mock_session, patch(
|
||
"application.streaming.message_journal.MessageEventsRepository"
|
||
) as mock_repo_cls, patch(
|
||
"application.streaming.message_journal.Topic"
|
||
) as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo_cls.return_value.record = MagicMock()
|
||
mock_topic = MagicMock()
|
||
mock_topic.publish.side_effect = RuntimeError("redis down")
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
# Must not raise.
|
||
result = record_event("msg-1", 0, "answer", {"answer": "y"})
|
||
assert result is True # Journal still committed.
|
||
|
||
def test_payload_none_treated_as_empty_dict(self):
|
||
with patch(
|
||
"application.streaming.message_journal.db_session"
|
||
) as mock_session, patch(
|
||
"application.streaming.message_journal.MessageEventsRepository"
|
||
) as mock_repo_cls, patch(
|
||
"application.streaming.message_journal.Topic"
|
||
) as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo = MagicMock()
|
||
mock_repo_cls.return_value = mock_repo
|
||
mock_topic_cls.return_value = MagicMock()
|
||
|
||
record_event("msg-1", 0, "end", None)
|
||
mock_repo.record.assert_called_once_with("msg-1", 0, "end", {})
|
||
|
||
def test_payload_non_dict_rejected_at_gate(self):
|
||
"""Contract: payload must be a dict (or None). Lists, strings,
|
||
ints, and other shapes are rejected without writing or
|
||
publishing.
|
||
|
||
Background: the live path (``base.py::_emit``) and the replay
|
||
path (``event_replay``) previously reconstructed non-dicts
|
||
differently — ``{"value": payload}`` live vs.
|
||
``{"type": event_type}`` on replay — so a reconnecting client
|
||
would receive a different envelope than the one originally
|
||
streamed. Rejecting at this gate keeps the two paths
|
||
byte-identical.
|
||
"""
|
||
with patch(
|
||
"application.streaming.message_journal.db_session"
|
||
) as mock_session, patch(
|
||
"application.streaming.message_journal.MessageEventsRepository"
|
||
) as mock_repo_cls, patch(
|
||
"application.streaming.message_journal.Topic"
|
||
) as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo = MagicMock()
|
||
mock_repo_cls.return_value = mock_repo
|
||
mock_topic = MagicMock()
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
result = record_event("msg-1", 0, "end", ["unexpected", "list"])
|
||
|
||
assert result is False
|
||
mock_repo.record.assert_not_called()
|
||
mock_topic.publish.assert_not_called()
|
||
|
||
def test_integrity_error_retries_with_seq_plus_one(self):
|
||
"""Composite-PK collision on (message_id, sequence_no) is
|
||
recovered by one retry against ``latest_sequence_no + 1`` —
|
||
the most likely cause is a stale seq seed on a continuation
|
||
retry, where the route read MAX(seq) from a separate
|
||
connection before another writer committed past it.
|
||
|
||
On success the live pubsub publish uses the retried seq so
|
||
the journal row and the live frame agree, even though the
|
||
caller's original POST stream still carries the pre-retry id.
|
||
"""
|
||
from sqlalchemy.exc import IntegrityError
|
||
|
||
# Two repo instances: one per ``with db_session()`` block.
|
||
# repo_first.record raises IntegrityError; the readonly
|
||
# session returns latest=7; repo_retry.record succeeds at 8.
|
||
repo_first = MagicMock(name="repo_first")
|
||
repo_first.record.side_effect = IntegrityError("stmt", {}, Exception())
|
||
repo_readonly = MagicMock(name="repo_readonly")
|
||
repo_readonly.latest_sequence_no.return_value = 7
|
||
repo_retry = MagicMock(name="repo_retry")
|
||
|
||
repo_instances = iter([repo_first, repo_readonly, repo_retry])
|
||
|
||
with patch(
|
||
"application.streaming.message_journal.db_session"
|
||
) as mock_session, patch(
|
||
"application.streaming.message_journal.db_readonly"
|
||
) as mock_readonly, patch(
|
||
"application.streaming.message_journal.MessageEventsRepository",
|
||
side_effect=lambda conn: next(repo_instances),
|
||
), patch(
|
||
"application.streaming.message_journal.Topic"
|
||
) as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_readonly.return_value.__enter__.return_value = MagicMock()
|
||
mock_topic = MagicMock()
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
result = record_event("msg-1", 3, "answer", {"text": "hi"})
|
||
|
||
assert result is True
|
||
# First INSERT attempted at the caller's seq=3.
|
||
repo_first.record.assert_called_once_with(
|
||
"msg-1", 3, "answer", {"text": "hi"}
|
||
)
|
||
# Latest probed for the retry seed.
|
||
repo_readonly.latest_sequence_no.assert_called_once_with("msg-1")
|
||
# Retry INSERT at latest+1 = 8.
|
||
repo_retry.record.assert_called_once_with(
|
||
"msg-1", 8, "answer", {"text": "hi"}
|
||
)
|
||
# Live publish uses the materialised seq so the wire and
|
||
# the journal row stay in lockstep on the retry path.
|
||
mock_topic.publish.assert_called_once()
|
||
wire = json.loads(mock_topic.publish.call_args[0][0])
|
||
assert wire["sequence_no"] == 8
|
||
|
||
def test_integrity_error_retry_failure_drops_silently(self):
|
||
"""If the retry collides again (truly concurrent writers in
|
||
lockstep) the journal write is dropped but the function still
|
||
returns ``False`` without raising — the streaming loop must
|
||
not be killed by a journal hiccup.
|
||
"""
|
||
from sqlalchemy.exc import IntegrityError
|
||
|
||
repo_first = MagicMock(name="repo_first")
|
||
repo_first.record.side_effect = IntegrityError("stmt", {}, Exception())
|
||
repo_readonly = MagicMock(name="repo_readonly")
|
||
repo_readonly.latest_sequence_no.return_value = 3
|
||
repo_retry = MagicMock(name="repo_retry")
|
||
repo_retry.record.side_effect = IntegrityError("stmt", {}, Exception())
|
||
|
||
repo_instances = iter([repo_first, repo_readonly, repo_retry])
|
||
|
||
with patch(
|
||
"application.streaming.message_journal.db_session"
|
||
) as mock_session, patch(
|
||
"application.streaming.message_journal.db_readonly"
|
||
) as mock_readonly, patch(
|
||
"application.streaming.message_journal.MessageEventsRepository",
|
||
side_effect=lambda conn: next(repo_instances),
|
||
), patch(
|
||
"application.streaming.message_journal.Topic"
|
||
) as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_readonly.return_value.__enter__.return_value = MagicMock()
|
||
mock_topic = MagicMock()
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
result = record_event("msg-1", 0, "answer", {"text": "hi"})
|
||
|
||
assert result is False
|
||
# Both INSERT attempts fired; the second raised too.
|
||
assert repo_first.record.call_count == 1
|
||
assert repo_retry.record.call_count == 1
|
||
# Live publish still fires on the materialised seq path —
|
||
# subscribers downgrade to keepalives if they were waiting
|
||
# on this event, which is correct since the journal can't
|
||
# serve it on a future reconnect anyway.
|
||
mock_topic.publish.assert_called_once()
|
||
|
||
def test_foreign_key_violation_drops_without_seq_retry(self):
|
||
"""A 23503 (foreign-key violation) means the
|
||
``conversation_messages`` parent row was never committed — a
|
||
``sequence_no`` retry can never land the row, so the write is
|
||
dropped immediately without probing ``latest_sequence_no`` or
|
||
re-attempting. The live publish still fires (best-effort tail).
|
||
"""
|
||
repo_first = MagicMock(name="repo_first")
|
||
repo_first.record.side_effect = _integrity_error("23503")
|
||
repo_readonly = MagicMock(name="repo_readonly")
|
||
repo_retry = MagicMock(name="repo_retry")
|
||
|
||
repo_instances = iter([repo_first, repo_readonly, repo_retry])
|
||
|
||
with patch(
|
||
"application.streaming.message_journal.db_session"
|
||
) as mock_session, patch(
|
||
"application.streaming.message_journal.db_readonly"
|
||
) as mock_readonly, patch(
|
||
"application.streaming.message_journal.MessageEventsRepository",
|
||
side_effect=lambda conn: next(repo_instances),
|
||
), patch(
|
||
"application.streaming.message_journal.Topic"
|
||
) as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_readonly.return_value.__enter__.return_value = MagicMock()
|
||
mock_topic = MagicMock()
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
result = record_event("msg-1", 3, "answer", {"text": "hi"})
|
||
|
||
assert result is False
|
||
repo_first.record.assert_called_once_with(
|
||
"msg-1", 3, "answer", {"text": "hi"}
|
||
)
|
||
# No seq+1 retry: the readonly probe and retry insert never run.
|
||
mock_readonly.assert_not_called()
|
||
repo_readonly.latest_sequence_no.assert_not_called()
|
||
repo_retry.record.assert_not_called()
|
||
# Live tail still gets the frame at the original seq.
|
||
mock_topic.publish.assert_called_once()
|
||
wire = json.loads(mock_topic.publish.call_args[0][0])
|
||
assert wire["sequence_no"] == 3
|
||
|
||
def test_pk_collision_still_retries_with_seq_plus_one(self):
|
||
"""An explicit 23505 (unique/PK collision) keeps the existing
|
||
recover-by-retry behavior: probe ``latest_sequence_no`` and
|
||
rewrite at latest+1. Regression guard pinning the FK-vs-PK split.
|
||
"""
|
||
repo_first = MagicMock(name="repo_first")
|
||
repo_first.record.side_effect = _integrity_error("23505")
|
||
repo_readonly = MagicMock(name="repo_readonly")
|
||
repo_readonly.latest_sequence_no.return_value = 7
|
||
repo_retry = MagicMock(name="repo_retry")
|
||
|
||
repo_instances = iter([repo_first, repo_readonly, repo_retry])
|
||
|
||
with patch(
|
||
"application.streaming.message_journal.db_session"
|
||
) as mock_session, patch(
|
||
"application.streaming.message_journal.db_readonly"
|
||
) as mock_readonly, patch(
|
||
"application.streaming.message_journal.MessageEventsRepository",
|
||
side_effect=lambda conn: next(repo_instances),
|
||
), patch(
|
||
"application.streaming.message_journal.Topic"
|
||
) as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_readonly.return_value.__enter__.return_value = MagicMock()
|
||
mock_topic_cls.return_value = MagicMock()
|
||
|
||
result = record_event("msg-1", 3, "answer", {"text": "hi"})
|
||
|
||
assert result is True
|
||
repo_readonly.latest_sequence_no.assert_called_once_with("msg-1")
|
||
repo_retry.record.assert_called_once_with(
|
||
"msg-1", 8, "answer", {"text": "hi"}
|
||
)
|
||
|
||
def test_payload_with_null_byte_stripped_before_insert(self):
|
||
"""``\\x00`` in a string value would crash JSONB INSERT — strip
|
||
at the journal boundary so the row lands and reconnecting
|
||
clients see the chunk.
|
||
"""
|
||
with patch(
|
||
"application.streaming.message_journal.db_session"
|
||
) as mock_session, patch(
|
||
"application.streaming.message_journal.MessageEventsRepository"
|
||
) as mock_repo_cls, patch(
|
||
"application.streaming.message_journal.Topic"
|
||
) as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo = MagicMock()
|
||
mock_repo_cls.return_value = mock_repo
|
||
mock_topic_cls.return_value = MagicMock()
|
||
|
||
record_event(
|
||
"msg-1", 0, "answer", {"answer": "hel\x00lo", "type": "answer"}
|
||
)
|
||
|
||
mock_repo.record.assert_called_once_with(
|
||
"msg-1", 0, "answer", {"answer": "hello", "type": "answer"}
|
||
)
|
||
|
||
def test_payload_with_nested_null_bytes_stripped_recursively(self):
|
||
"""Strip walks nested dicts and lists — and string keys, to be
|
||
safe against an LLM emitting structured data with NULs.
|
||
"""
|
||
with patch(
|
||
"application.streaming.message_journal.db_session"
|
||
) as mock_session, patch(
|
||
"application.streaming.message_journal.MessageEventsRepository"
|
||
) as mock_repo_cls, patch(
|
||
"application.streaming.message_journal.Topic"
|
||
) as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo = MagicMock()
|
||
mock_repo_cls.return_value = mock_repo
|
||
mock_topic_cls.return_value = MagicMock()
|
||
|
||
record_event(
|
||
"msg-1",
|
||
0,
|
||
"answer",
|
||
{
|
||
"outer": {"inner\x00key": "a\x00b"},
|
||
"items": ["x\x00y", {"k": "v\x00"}],
|
||
"clean": "ok",
|
||
},
|
||
)
|
||
|
||
mock_repo.record.assert_called_once_with(
|
||
"msg-1",
|
||
0,
|
||
"answer",
|
||
{
|
||
"outer": {"innerkey": "ab"},
|
||
"items": ["xy", {"k": "v"}],
|
||
"clean": "ok",
|
||
},
|
||
)
|
||
|
||
def test_payload_without_null_bytes_unchanged(self):
|
||
"""Common-case: no NUL, no allocation churn — payload passed
|
||
through identically.
|
||
"""
|
||
with patch(
|
||
"application.streaming.message_journal.db_session"
|
||
) as mock_session, patch(
|
||
"application.streaming.message_journal.MessageEventsRepository"
|
||
) as mock_repo_cls, patch(
|
||
"application.streaming.message_journal.Topic"
|
||
) as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo = MagicMock()
|
||
mock_repo_cls.return_value = mock_repo
|
||
mock_topic_cls.return_value = MagicMock()
|
||
|
||
payload = {"a": "1", "nested": {"b": [1, 2, "three"]}}
|
||
record_event("msg-1", 0, "answer", payload)
|
||
|
||
mock_repo.record.assert_called_once_with(
|
||
"msg-1", 0, "answer", payload
|
||
)
|
||
|
||
|
||
@pytest.mark.unit
|
||
class TestBatchedJournalWriter:
|
||
"""``BatchedJournalWriter`` amortizes per-emit PG commits.
|
||
|
||
Tests cover the four flush triggers (size, time, lifecycle,
|
||
explicit), the bulk→per-row IntegrityError fallback, and the
|
||
live-publish-stays-synchronous invariant.
|
||
"""
|
||
|
||
def _patch_io(self):
|
||
"""Standard patch set: stub PG sessions and Redis pubsub.
|
||
|
||
Returns a tuple ``(session_cm, readonly_cm, repo_factory, topic)``.
|
||
The repo_factory is a list — each ``MessageEventsRepository(...)``
|
||
call inside the writer pops one element off the front. Tests
|
||
push fakes onto it before triggering the path that needs them.
|
||
"""
|
||
from unittest.mock import patch as _patch
|
||
|
||
session_cm = _patch("application.streaming.message_journal.db_session")
|
||
readonly_cm = _patch(
|
||
"application.streaming.message_journal.db_readonly"
|
||
)
|
||
repo_cls = _patch(
|
||
"application.streaming.message_journal.MessageEventsRepository"
|
||
)
|
||
topic_cls = _patch("application.streaming.message_journal.Topic")
|
||
return session_cm, readonly_cm, repo_cls, topic_cls
|
||
|
||
def test_size_trigger_flushes_at_batch_size(self):
|
||
"""Buffer reaches ``batch_size`` → one bulk_record call covers
|
||
all rows; publishes fire after the bulk commit, one per row.
|
||
"""
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo = MagicMock()
|
||
mock_repo_cls.return_value = mock_repo
|
||
mock_topic = MagicMock()
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
writer = BatchedJournalWriter(
|
||
"msg-1", batch_size=3, batch_interval_ms=10_000
|
||
)
|
||
# First two records: buffered, no bulk_record yet, no publish.
|
||
writer.record(0, "answer", {"text": "a"})
|
||
writer.record(1, "answer", {"text": "b"})
|
||
mock_repo.bulk_record.assert_not_called()
|
||
assert mock_topic.publish.call_count == 0
|
||
# Third record: triggers a size-based flush.
|
||
writer.record(2, "answer", {"text": "c"})
|
||
mock_repo.bulk_record.assert_called_once()
|
||
args = mock_repo.bulk_record.call_args
|
||
assert args.args[0] == "msg-1"
|
||
buffered = args.args[1]
|
||
assert len(buffered) == 3
|
||
assert [seq for seq, _, _ in buffered] == [0, 1, 2]
|
||
# Publishes fire after the bulk commit, one per row in order.
|
||
assert mock_topic.publish.call_count == 3
|
||
published_seqs = [
|
||
json.loads(call.args[0])["sequence_no"]
|
||
for call in mock_topic.publish.call_args_list
|
||
]
|
||
assert published_seqs == [0, 1, 2]
|
||
|
||
def test_time_trigger_flushes_after_interval(self):
|
||
"""When the elapsed time since the last flush exceeds
|
||
``batch_interval_ms``, the next ``record()`` flushes — even if
|
||
the buffer is well below ``batch_size``. Drives reconnect
|
||
visibility for slow producers.
|
||
"""
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls, patch(
|
||
"application.streaming.message_journal.time.monotonic"
|
||
) as mock_mono:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo = MagicMock()
|
||
mock_repo_cls.return_value = mock_repo
|
||
mock_topic_cls.return_value = MagicMock()
|
||
|
||
# First call from __init__ at t=0; subsequent calls drive
|
||
# the time-trigger check inside ``_should_flush``.
|
||
# t=0.000 → __init__ snapshots last_flush
|
||
# t=0.001 → record(0) — 1ms elapsed → no flush
|
||
# t=0.001 → after record, last_flush snapshot when flush called (n/a)
|
||
# t=0.200 → record(1) — 200ms elapsed → flush
|
||
# t=0.200 → flush updates last_flush
|
||
mock_mono.side_effect = iter([0.000, 0.001, 0.200, 0.200])
|
||
|
||
writer = BatchedJournalWriter(
|
||
"msg-1", batch_size=100, batch_interval_ms=100
|
||
)
|
||
writer.record(0, "answer", {"text": "a"})
|
||
mock_repo.bulk_record.assert_not_called()
|
||
writer.record(1, "answer", {"text": "b"})
|
||
mock_repo.bulk_record.assert_called_once()
|
||
|
||
def test_close_drains_remaining_buffer(self):
|
||
"""``close()`` is the lifecycle flush — at end of stream every
|
||
buffered event must commit before the writer goes silent.
|
||
Idempotent so it's safe in multiple finally clauses.
|
||
"""
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo = MagicMock()
|
||
mock_repo_cls.return_value = mock_repo
|
||
mock_topic_cls.return_value = MagicMock()
|
||
|
||
writer = BatchedJournalWriter(
|
||
"msg-1", batch_size=100, batch_interval_ms=100_000
|
||
)
|
||
writer.record(0, "answer", {"text": "a"})
|
||
writer.record(1, "end", {"type": "end"})
|
||
mock_repo.bulk_record.assert_not_called()
|
||
|
||
writer.close()
|
||
mock_repo.bulk_record.assert_called_once()
|
||
assert len(mock_repo.bulk_record.call_args.args[1]) == 2
|
||
|
||
# Second close is a no-op.
|
||
writer.close()
|
||
assert mock_repo.bulk_record.call_count == 1
|
||
|
||
def test_record_after_close_returns_false(self):
|
||
"""Writing past ``close()`` must not silently land in a stale
|
||
buffer — return False so the caller's flow surfaces the bug.
|
||
"""
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo_cls.return_value = MagicMock()
|
||
mock_topic_cls.return_value = MagicMock()
|
||
|
||
writer = BatchedJournalWriter("msg-1", batch_size=100)
|
||
writer.close()
|
||
assert writer.record(0, "answer", {"text": "a"}) is False
|
||
|
||
def test_record_rejects_non_dict_payload(self):
|
||
"""Same contract as ``record_event``: non-dict payloads are
|
||
rejected to keep live and replay paths byte-identical.
|
||
"""
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo_cls.return_value = MagicMock()
|
||
mock_topic = MagicMock()
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
writer = BatchedJournalWriter("msg-1")
|
||
assert writer.record(0, "answer", ["bad", "list"]) is False
|
||
# Nothing publishes either — the gate fires before the wire.
|
||
mock_topic.publish.assert_not_called()
|
||
|
||
def test_bulk_collision_falls_back_to_per_row(self):
|
||
"""Bulk INSERT fails with IntegrityError → writer retries each
|
||
row individually via the legacy ``record_event`` retry path,
|
||
so a single colliding seq doesn't drop the whole batch.
|
||
"""
|
||
from sqlalchemy.exc import IntegrityError
|
||
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm as mock_readonly, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_readonly.return_value.__enter__.return_value = MagicMock()
|
||
|
||
# First call: bulk_record raises. Subsequent calls: per-row
|
||
# record succeeds. Use a single fake whose first
|
||
# ``bulk_record`` call raises and whose ``record`` succeeds.
|
||
bulk_repo = MagicMock(name="bulk_repo")
|
||
bulk_repo.bulk_record.side_effect = IntegrityError(
|
||
"stmt", {}, Exception()
|
||
)
|
||
per_row_repo = MagicMock(name="per_row_repo")
|
||
# MessageEventsRepository(conn) is called once per session
|
||
# opened. Bulk path = 1; per-row fallback = 2 rows × 1 each.
|
||
mock_repo_cls.side_effect = [bulk_repo, per_row_repo, per_row_repo]
|
||
|
||
mock_topic_cls.return_value = MagicMock()
|
||
|
||
writer = BatchedJournalWriter("msg-1", batch_size=2)
|
||
writer.record(0, "answer", {"text": "a"})
|
||
writer.record(1, "answer", {"text": "b"})
|
||
|
||
bulk_repo.bulk_record.assert_called_once()
|
||
# Per-row fallback wrote each row in its own session.
|
||
assert per_row_repo.record.call_count == 2
|
||
assert per_row_repo.record.call_args_list[0].args[1] == 0
|
||
assert per_row_repo.record.call_args_list[1].args[1] == 1
|
||
|
||
def test_bulk_foreign_key_violation_drops_batch_without_per_row(self):
|
||
"""A 23503 (foreign-key violation) on the bulk INSERT means the
|
||
whole batch references a ``conversation_messages`` parent that
|
||
doesn't exist — every per-row retry would FK-fail too. The
|
||
batch is dropped once without the per-row fallback (no latest
|
||
probe, no re-attempts, no publish), and the buffer is cleared.
|
||
"""
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm as mock_readonly, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_readonly.return_value.__enter__.return_value = MagicMock()
|
||
|
||
bulk_repo = MagicMock(name="bulk_repo")
|
||
bulk_repo.bulk_record.side_effect = _integrity_error("23503")
|
||
per_row_repo = MagicMock(name="per_row_repo")
|
||
mock_repo_cls.side_effect = [bulk_repo, per_row_repo]
|
||
|
||
mock_topic = MagicMock()
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
writer = BatchedJournalWriter("msg-1", batch_size=2)
|
||
writer.record(0, "answer", {"text": "a"})
|
||
writer.record(1, "answer", {"text": "b"})
|
||
|
||
bulk_repo.bulk_record.assert_called_once()
|
||
# No per-row fallback: the parent is missing for the whole batch.
|
||
mock_readonly.assert_not_called()
|
||
per_row_repo.record.assert_not_called()
|
||
mock_topic.publish.assert_not_called()
|
||
# Buffer cleared so close() is a no-op and memory stays bounded.
|
||
assert writer._buffer == []
|
||
|
||
def test_bulk_collision_23505_still_falls_back_to_per_row(self):
|
||
"""An explicit 23505 (unique/PK collision) keeps the per-row
|
||
fallback — the parent exists, only a ``sequence_no`` collided.
|
||
Regression guard for the FK-vs-PK split at the bulk boundary.
|
||
"""
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm as mock_readonly, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_readonly.return_value.__enter__.return_value = MagicMock()
|
||
|
||
bulk_repo = MagicMock(name="bulk_repo")
|
||
bulk_repo.bulk_record.side_effect = _integrity_error("23505")
|
||
per_row_repo = MagicMock(name="per_row_repo")
|
||
mock_repo_cls.side_effect = [bulk_repo, per_row_repo, per_row_repo]
|
||
|
||
mock_topic_cls.return_value = MagicMock()
|
||
|
||
writer = BatchedJournalWriter("msg-1", batch_size=2)
|
||
writer.record(0, "answer", {"text": "a"})
|
||
writer.record(1, "answer", {"text": "b"})
|
||
|
||
bulk_repo.bulk_record.assert_called_once()
|
||
assert per_row_repo.record.call_count == 2
|
||
|
||
def test_flush_clears_buffer_even_on_total_failure(self):
|
||
"""A flush that fails with a non-IntegrityError exception must
|
||
still clear the buffer — leaving rows would grow memory
|
||
unbounded across the remainder of the stream. Degraded UX
|
||
(missing snapshot rows) beats runaway memory.
|
||
"""
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo = MagicMock()
|
||
mock_repo.bulk_record.side_effect = RuntimeError("PG gone")
|
||
mock_repo_cls.return_value = mock_repo
|
||
mock_topic_cls.return_value = MagicMock()
|
||
|
||
writer = BatchedJournalWriter("msg-1", batch_size=2)
|
||
writer.record(0, "answer", {"text": "a"})
|
||
writer.record(1, "answer", {"text": "b"})
|
||
# Bulk failed, but the buffer is cleared so a subsequent
|
||
# close() is a no-op and memory stays bounded.
|
||
assert writer._buffer == []
|
||
writer.close()
|
||
# Only the one flush we forced — no double-attempt on
|
||
# close after the buffer was already drained.
|
||
assert mock_repo.bulk_record.call_count == 1
|
||
|
||
def test_record_does_not_publish_synchronously(self):
|
||
"""Regression guard: ``record()`` must not publish before the
|
||
journal INSERT commits. Buffering an event with no flush yet
|
||
means zero pubsub frames on the wire — otherwise a reconnect
|
||
snapshot could miss a row that subscribers already received.
|
||
"""
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo_cls.return_value = MagicMock()
|
||
mock_topic = MagicMock()
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
writer = BatchedJournalWriter("msg-1", batch_size=100)
|
||
for i in range(5):
|
||
writer.record(i, "answer", {"text": str(i)})
|
||
mock_topic.publish.assert_not_called()
|
||
|
||
def test_flush_publishes_each_buffered_event_after_bulk_insert_in_order(
|
||
self,
|
||
):
|
||
"""Happy path: flush issues one bulk INSERT, then publishes each
|
||
buffered frame in the order it was recorded.
|
||
"""
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo = MagicMock()
|
||
mock_repo_cls.return_value = mock_repo
|
||
mock_topic = MagicMock()
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
writer = BatchedJournalWriter(
|
||
"msg-1", batch_size=100, batch_interval_ms=100_000
|
||
)
|
||
for i in range(4):
|
||
writer.record(i, "answer", {"text": str(i)})
|
||
mock_topic.publish.assert_not_called()
|
||
|
||
writer.flush()
|
||
mock_repo.bulk_record.assert_called_once()
|
||
assert mock_topic.publish.call_count == 4
|
||
published_seqs = [
|
||
json.loads(call.args[0])["sequence_no"]
|
||
for call in mock_topic.publish.call_args_list
|
||
]
|
||
assert published_seqs == [0, 1, 2, 3]
|
||
|
||
def test_flush_bulk_exception_drops_both_row_and_publish(self):
|
||
"""Non-IntegrityError from bulk_record → both the journal row
|
||
and the live publish are dropped together. Matches the
|
||
publish_user_event "drop live publish when durable write fails"
|
||
contract.
|
||
"""
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo = MagicMock()
|
||
mock_repo.bulk_record.side_effect = RuntimeError("PG gone")
|
||
mock_repo_cls.return_value = mock_repo
|
||
mock_topic = MagicMock()
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
writer = BatchedJournalWriter("msg-1", batch_size=2)
|
||
writer.record(0, "answer", {"text": "a"})
|
||
writer.record(1, "answer", {"text": "b"})
|
||
|
||
mock_repo.bulk_record.assert_called_once()
|
||
mock_topic.publish.assert_not_called()
|
||
|
||
def test_flush_per_row_publishes_on_commit_skips_on_persistent_drop(self):
|
||
"""IntegrityError-after-retry path: per-row fallback publishes
|
||
each row that lands in PG (including those rewritten at
|
||
latest+1), and skips publish for any row dropped after the
|
||
retry collision.
|
||
"""
|
||
from sqlalchemy.exc import IntegrityError
|
||
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm as mock_readonly, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_readonly.return_value.__enter__.return_value = MagicMock()
|
||
|
||
# Bulk attempt collides. Per-row: row 0 commits at its seq;
|
||
# row 1 collides, readonly probe returns latest=9, retry at
|
||
# 10 collides too → dropped without publish.
|
||
bulk_repo = MagicMock(name="bulk_repo")
|
||
bulk_repo.bulk_record.side_effect = IntegrityError(
|
||
"stmt", {}, Exception()
|
||
)
|
||
row0_repo = MagicMock(name="row0_repo")
|
||
row1_first_repo = MagicMock(name="row1_first_repo")
|
||
row1_first_repo.record.side_effect = IntegrityError(
|
||
"stmt", {}, Exception()
|
||
)
|
||
row1_readonly_repo = MagicMock(name="row1_readonly_repo")
|
||
row1_readonly_repo.latest_sequence_no.return_value = 9
|
||
row1_retry_repo = MagicMock(name="row1_retry_repo")
|
||
row1_retry_repo.record.side_effect = IntegrityError(
|
||
"stmt", {}, Exception()
|
||
)
|
||
mock_repo_cls.side_effect = [
|
||
bulk_repo,
|
||
row0_repo,
|
||
row1_first_repo,
|
||
row1_readonly_repo,
|
||
row1_retry_repo,
|
||
]
|
||
|
||
mock_topic = MagicMock()
|
||
mock_topic_cls.return_value = mock_topic
|
||
|
||
writer = BatchedJournalWriter("msg-1", batch_size=2)
|
||
writer.record(0, "answer", {"text": "a"})
|
||
writer.record(1, "answer", {"text": "b"})
|
||
|
||
# Only row 0 publishes — row 1 was dropped after the retry
|
||
# collision, so no live frame reaches the wire for it.
|
||
assert mock_topic.publish.call_count == 1
|
||
wire = json.loads(mock_topic.publish.call_args[0][0])
|
||
assert wire["sequence_no"] == 0
|
||
assert wire["payload"] == {"text": "a"}
|
||
|
||
def test_record_strips_null_bytes_before_buffer(self):
|
||
"""Null-byte sanitization applies at the writer boundary too —
|
||
otherwise a NUL-bearing chunk would land in the buffer, the
|
||
bulk INSERT would fail with ``DataError``, and the per-row
|
||
fallback would re-attempt the same broken row. Strip at the
|
||
gate so the buffered tuple is JSONB-safe.
|
||
"""
|
||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||
with session_cm as mock_session, readonly_cm, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||
mock_repo = MagicMock()
|
||
mock_repo_cls.return_value = mock_repo
|
||
mock_topic_cls.return_value = MagicMock()
|
||
|
||
writer = BatchedJournalWriter("msg-1", batch_size=2)
|
||
writer.record(0, "answer", {"answer": "hel\x00lo"})
|
||
writer.record(
|
||
1, "answer", {"nested": {"k\x00": ["a\x00b", "ok"]}}
|
||
)
|
||
|
||
mock_repo.bulk_record.assert_called_once()
|
||
buffered = mock_repo.bulk_record.call_args.args[1]
|
||
assert buffered[0][2] == {"answer": "hello"}
|
||
assert buffered[1][2] == {"nested": {"k": ["ab", "ok"]}}
|