Chunki z NUL wywalaly insert do postgresa (asyncpg CharacterNotInRepertoireError: invalid byte sequence for encoding "UTF8": 0x00) - 3 przypadki na mailach z 2007 w plastrze offset 50000 Etapu B. sanitize_surrogates tego nie lapie, bo NUL to poprawny code point, nie osierocony surogat. Strip dzieje sie zaraz po strip_quotes, czyli PRZED chunk_text i przed wywolaniem Ollamy - dzieki temu embedding liczy sie z dokladnie tego samego stringa, ktory trafia do document_chunk.text. Sanityzacja dopiero przy insercie zostawialaby wektor opisujacy tekst, ktorego DB nigdy nie zobaczyla. Skala widoczna w progress/summary: nul_bytes_stripped (ile znakow) oraz mails_nul_sanitized (ilu maili dotyczylo). Zadne z nich nie wchodzi do rownan balansu i nie wplywa na exit code - to normalizacja, nie blad. Ten sam strip w extract_threading (In-Reply-To / References): json.dumps zamienia NUL na escape u0000, ktory jsonb odrzuca tym samym bledem. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
859 lines
38 KiB
Python
859 lines
38 KiB
Python
"""Unit tests for the mail body ingest job — no DB, no real HTTP, no real Ollama.
|
|
|
|
Fixture .eml bytes are built inline (email.mime helpers / raw byte literals), matching the
|
|
convention in jobs/gmail-header-backfill/tests/test_backfill.py rather than separate fixture
|
|
files -- easier to see exactly what each test is asserting against."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from datetime import datetime, timezone
|
|
from email.mime.application import MIMEApplication
|
|
from email.mime.multipart import MIMEMultipart
|
|
from email.mime.text import MIMEText
|
|
|
|
import aiohttp
|
|
import pytest
|
|
|
|
from kb_retrieval.embed import EmbeddingDimensionError
|
|
from mail_body_ingest.ingest import (
|
|
_CHUNK_INSERT_SQL,
|
|
EmbedBackendUnavailableError,
|
|
_decode_jsonb,
|
|
_embed_ms_per_chunk,
|
|
_env_num,
|
|
_has_threading,
|
|
_new_stats,
|
|
build_prefix,
|
|
extract_body,
|
|
extract_threading,
|
|
fetch_envelopes,
|
|
fetch_existing_chunk_keys,
|
|
html_to_text,
|
|
insert_chunk,
|
|
is_newsletter,
|
|
parse_message,
|
|
run,
|
|
strip_quotes,
|
|
)
|
|
|
|
|
|
def _plain_eml(headers: dict, body: str) -> bytes:
|
|
msg = MIMEText(body, "plain", "utf-8")
|
|
for k, v in headers.items():
|
|
msg[k] = v
|
|
return msg.as_bytes()
|
|
|
|
|
|
def _html_eml(headers: dict, html_body: str) -> bytes:
|
|
msg = MIMEText(html_body, "html", "utf-8")
|
|
for k, v in headers.items():
|
|
msg[k] = v
|
|
return msg.as_bytes()
|
|
|
|
|
|
def _multipart_with_attachment_eml(headers: dict, body: str) -> bytes:
|
|
msg = MIMEMultipart()
|
|
for k, v in headers.items():
|
|
msg[k] = v
|
|
msg.attach(MIMEText(body, "plain", "utf-8"))
|
|
attachment = MIMEApplication(b"%PDF-fake", Name="doc.pdf")
|
|
attachment["Content-Disposition"] = 'attachment; filename="doc.pdf"'
|
|
msg.attach(attachment)
|
|
return msg.as_bytes()
|
|
|
|
|
|
# Pattern 1 from the gmail-header-backfill diagnosis (2026-07-14): RFC 2047 encoded-word whose
|
|
# decoded text contains a newline in a display name -- policy.default raises ValueError.
|
|
_CRLF_ENCODED_WORD_EML = (
|
|
b"From: Rekrutacja z =?utf-8?Q?ExampleCorp=0A?= <mailing@example.pl>\r\n"
|
|
b"To: Jan Kowalski <jan@example.com>\r\n"
|
|
b"Subject: Nowe oferty\r\n"
|
|
b"Date: Mon, 15 Sep 2014 13:24:18 +0000\r\n"
|
|
b"\r\nbody"
|
|
)
|
|
|
|
|
|
class TestStripQuotes:
|
|
def test_no_marker_no_quotes_returns_unchanged(self):
|
|
clean, stripped = strip_quotes("Hello,\n\nJust checking in.\n\nRegards,\nA")
|
|
assert clean == "Hello,\n\nJust checking in.\n\nRegards,\nA"
|
|
assert stripped == 0
|
|
|
|
def test_gmail_wrote_marker_truncates_everything_after(self):
|
|
text = "Sure, sounds good.\n\nOn Tue, 10 Jun 2025 at 12:00, Alice <a@b.com> wrote:\n> old text\n> more old"
|
|
clean, stripped = strip_quotes(text)
|
|
assert clean == "Sure, sounds good."
|
|
assert stripped > 0
|
|
|
|
def test_polish_dnia_napisal_marker(self):
|
|
text = "OK, dzieki.\n\nDnia 10 czerwca 2025 Jan napisał(a):\n> stara tresc"
|
|
clean, _ = strip_quotes(text)
|
|
assert clean == "OK, dzieki."
|
|
|
|
def test_polish_w_dniu_pisze_marker(self):
|
|
text = "Dobrze.\n\nW dniu 10.06.2025 Anna pisze:\n> cos tam"
|
|
clean, _ = strip_quotes(text)
|
|
assert clean == "Dobrze."
|
|
|
|
def test_outlook_original_message_marker(self):
|
|
text = "Ok, dzieki za info.\n\n-----Original Message-----\nFrom: bob@example.com\nold content"
|
|
clean, _ = strip_quotes(text)
|
|
assert clean == "Ok, dzieki za info."
|
|
|
|
def test_outlook_underscore_separator(self):
|
|
text = "Widziane, dzieki.\n\n________________________________\nFrom: c@example.com\nstuff"
|
|
clean, _ = strip_quotes(text)
|
|
assert clean == "Widziane, dzieki."
|
|
|
|
def test_bare_quote_lines_without_marker_are_dropped(self):
|
|
text = "My reply here.\n> quoted line one\n> quoted line two\nNot quoted trailing line"
|
|
clean, stripped = strip_quotes(text)
|
|
assert ">" not in clean
|
|
assert "My reply here." in clean
|
|
assert stripped > 0
|
|
|
|
def test_chars_stripped_counts_removed_characters(self):
|
|
text = "short reply\n\nOn 1 Jan 2020 X wrote:\n> a very long quoted paragraph indeed"
|
|
_clean, stripped = strip_quotes(text)
|
|
assert stripped == len(text) - len("short reply")
|
|
|
|
|
|
class TestHtmlToText:
|
|
def test_skips_style_and_script(self):
|
|
html = "<html><head><style>.x{color:red}</style></head><body><script>evil()</script><p>Real text</p></body></html>"
|
|
assert "color:red" not in html_to_text(html)
|
|
assert "evil()" not in html_to_text(html)
|
|
assert "Real text" in html_to_text(html)
|
|
|
|
def test_skips_blockquote_subtree(self):
|
|
html = "<div>Reply text</div><blockquote><div>quoted old text</div></blockquote>"
|
|
text = html_to_text(html)
|
|
assert "Reply text" in text
|
|
assert "quoted old text" not in text
|
|
|
|
def test_skips_gmail_quote_div(self):
|
|
html = (
|
|
'<div>My reply</div>'
|
|
'<div class="gmail_quote">On Tue, X wrote:<div>nested quoted content</div></div>'
|
|
)
|
|
text = html_to_text(html)
|
|
assert "My reply" in text
|
|
assert "nested quoted content" not in text
|
|
|
|
def test_plain_div_without_gmail_quote_class_is_not_skipped(self):
|
|
html = '<div class="not_a_quote">Still visible</div>'
|
|
assert "Still visible" in html_to_text(html)
|
|
|
|
def test_unclosed_void_elements_in_head_do_not_blank_the_rest_of_the_document(self):
|
|
# Regression: a naive LIFO stack pops the wrong "frame" for the literal </head> when
|
|
# <meta> tags inside <head> are written without a self-closing slash (the overwhelming
|
|
# majority of real-world email HTML) -- confirmed live against an Etap A pilot mail
|
|
# where this bug silently blanked the entire body to ''.
|
|
html = (
|
|
"<!DOCTYPE html><html><head>"
|
|
'<meta charset="utf-8"><meta name="viewport" content="width=device-width">'
|
|
"<title>promo</title>"
|
|
"</head><body><p>Kod rabatowy dla Ciebie!</p></body></html>"
|
|
)
|
|
assert "Kod rabatowy dla Ciebie!" in html_to_text(html)
|
|
|
|
def test_br_and_hr_void_elements_still_produce_line_breaks(self):
|
|
html = "<p>line one<br>line two</p><hr><p>line three</p>"
|
|
text = html_to_text(html)
|
|
assert "line one" in text and "line two" in text and "line three" in text
|
|
|
|
def test_stray_unmatched_closing_tag_is_ignored_not_fatal(self):
|
|
html = "<p>hello</blockquote><p>world</p>"
|
|
text = html_to_text(html)
|
|
assert "hello" in text and "world" in text
|
|
|
|
|
|
class TestIsNewsletter:
|
|
def test_list_unsubscribe_present(self):
|
|
msg = parse_message(_plain_eml({"From": "a@b.com", "List-Unsubscribe": "<mailto:x@y.com>"}, "hi"))
|
|
assert is_newsletter(msg) is True
|
|
|
|
def test_list_id_present(self):
|
|
msg = parse_message(_plain_eml({"From": "a@b.com", "List-Id": "newsletter.example.com"}, "hi"))
|
|
assert is_newsletter(msg) is True
|
|
|
|
def test_precedence_bulk(self):
|
|
msg = parse_message(_plain_eml({"From": "a@b.com", "Precedence": "bulk"}, "hi"))
|
|
assert is_newsletter(msg) is True
|
|
|
|
def test_precedence_list_case_insensitive(self):
|
|
msg = parse_message(_plain_eml({"From": "a@b.com", "Precedence": "LIST"}, "hi"))
|
|
assert is_newsletter(msg) is True
|
|
|
|
def test_ordinary_mail_is_not_newsletter(self):
|
|
msg = parse_message(_plain_eml({"From": "a@b.com", "Subject": "hi"}, "hi"))
|
|
assert is_newsletter(msg) is False
|
|
|
|
|
|
class TestExtractThreading:
|
|
def test_in_reply_to_and_references_stripped_of_brackets(self):
|
|
msg = parse_message(_plain_eml(
|
|
{"From": "a@b.com", "In-Reply-To": "<msg1@x.com>",
|
|
"References": "<msg0@x.com> <msg1@x.com>"}, "hi",
|
|
))
|
|
threading = extract_threading(msg)
|
|
assert threading == {
|
|
"type": "threading", "in_reply_to": "msg1@x.com",
|
|
"references": ["msg0@x.com", "msg1@x.com"],
|
|
}
|
|
|
|
def test_missing_headers_yield_none_and_empty_list(self):
|
|
msg = parse_message(_plain_eml({"From": "a@b.com"}, "hi"))
|
|
threading = extract_threading(msg)
|
|
assert threading == {"type": "threading", "in_reply_to": None, "references": []}
|
|
|
|
|
|
class TestBuildPrefix:
|
|
def test_full_headers(self):
|
|
entities = [{"type": "headers", "subject": "Twoja polisa", "from": {"name": "PZU", "address": "x@pzu.pl"}}]
|
|
ts = datetime(2025, 8, 3, tzinfo=timezone.utc)
|
|
assert build_prefix(entities, ts) == "Temat: Twoja polisa | Od: PZU | Data: 2025-08-03"
|
|
|
|
def test_missing_subject_and_from_fallback(self):
|
|
entities = [{"type": "headers", "subject": None, "from": None}]
|
|
ts = datetime(2025, 1, 1, tzinfo=timezone.utc)
|
|
prefix = build_prefix(entities, ts)
|
|
assert "(brak tematu)" in prefix
|
|
assert "Od: ?" in prefix
|
|
|
|
def test_from_without_name_uses_address(self):
|
|
entities = [{"type": "headers", "subject": "s", "from": {"name": None, "address": "a@b.com"}}]
|
|
ts = datetime(2025, 1, 1, tzinfo=timezone.utc)
|
|
assert "Od: a@b.com" in build_prefix(entities, ts)
|
|
|
|
def test_no_headers_entity_at_all(self):
|
|
ts = datetime(2025, 1, 1, tzinfo=timezone.utc)
|
|
prefix = build_prefix([], ts)
|
|
assert "(brak tematu)" in prefix
|
|
assert "Od: ?" in prefix
|
|
|
|
|
|
class TestExtractBody:
|
|
def test_plain_text_preferred(self):
|
|
msg = parse_message(_plain_eml({"From": "a@b.com"}, "hello plain body"))
|
|
assert extract_body(msg) == "hello plain body"
|
|
|
|
def test_html_only_falls_back_to_html_to_text(self):
|
|
msg = parse_message(_html_eml({"From": "a@b.com"}, "<p>hello <b>html</b> body</p>"))
|
|
body = extract_body(msg)
|
|
assert "hello" in body and "html" in body and "body" in body
|
|
assert "<" not in body
|
|
|
|
def test_attachment_only_yields_empty_body(self):
|
|
msg = MIMEApplication(b"%PDF-fake", Name="doc.pdf")
|
|
msg["From"] = "a@b.com"
|
|
assert extract_body(msg) == ""
|
|
|
|
def test_multipart_with_attachment_extracts_only_the_text_part(self):
|
|
msg = parse_message(_multipart_with_attachment_eml({"From": "a@b.com"}, "the actual message"))
|
|
assert extract_body(msg) == "the actual message"
|
|
|
|
|
|
class TestParseMessage:
|
|
def test_typed_parse_succeeds_for_ordinary_mail(self):
|
|
msg = parse_message(_plain_eml({"From": "a@b.com", "Subject": "hi"}, "body"))
|
|
assert msg.get("Subject") == "hi"
|
|
|
|
def test_falls_back_to_compat32_on_typed_parse_failure(self):
|
|
# This exact byte pattern raises ValueError under policy.default (see backfill's
|
|
# diagnosis) -- parse_message must still return a usable Message via the fallback.
|
|
msg = parse_message(_CRLF_ENCODED_WORD_EML)
|
|
assert msg.get("Subject") is not None
|
|
|
|
|
|
class TestDecodeJsonbAndThreadingHelpers:
|
|
def test_decode_jsonb_handles_str_and_object(self):
|
|
assert _decode_jsonb('[{"type": "headers"}]') == [{"type": "headers"}]
|
|
assert _decode_jsonb([{"type": "headers"}]) == [{"type": "headers"}]
|
|
assert _decode_jsonb(None) is None
|
|
|
|
def test_has_threading(self):
|
|
assert _has_threading([{"type": "threading"}]) is True
|
|
assert _has_threading([{"type": "headers"}]) is False
|
|
assert _has_threading([]) is False
|
|
|
|
|
|
class TestInsertSql:
|
|
def test_on_conflict_target_includes_model(self):
|
|
assert "ON CONFLICT (envelope_id, chunk_index, model)" in _CHUNK_INSERT_SQL
|
|
|
|
|
|
class TestInsertChunk:
|
|
async def test_newsletter_chunk_inserts_null_embedding_with_reason(self):
|
|
conn = _FakeConn()
|
|
await insert_chunk(conn, "msgid@x", 0, "junk text", None, "bge-m3", excluded_reason="newsletter")
|
|
assert conn.execute_calls == [("msgid@x", 0, "junk text", None, "bge-m3", "newsletter")]
|
|
|
|
async def test_normal_chunk_inserts_vector_with_no_reason(self):
|
|
conn = _FakeConn()
|
|
await insert_chunk(conn, "msgid@x", 0, "real text", [0.1, 0.2], "bge-m3")
|
|
assert conn.execute_calls == [("msgid@x", 0, "real text", "[0.1,0.2]", "bge-m3", None)]
|
|
|
|
|
|
class TestFetchHelpers:
|
|
async def test_fetch_envelopes_no_filters(self):
|
|
conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")])
|
|
rows = await fetch_envelopes(conn, since=None, limit=None, offset=None)
|
|
assert rows == [_env_row("m1@x", "hi")]
|
|
query, params = conn.queries[-1]
|
|
assert "ORDER BY id" in query
|
|
assert "LIMIT" not in query and "OFFSET" not in query
|
|
|
|
async def test_fetch_envelopes_with_since_limit_offset(self):
|
|
conn = _FakeConn(envelopes=[])
|
|
await fetch_envelopes(conn, since=datetime(2025, 7, 1, tzinfo=timezone.utc), limit=10, offset=5)
|
|
query, params = conn.queries[-1]
|
|
assert "ts >= $1" in query
|
|
assert "LIMIT $2" in query
|
|
assert "OFFSET $3" in query
|
|
assert params == (datetime(2025, 7, 1, tzinfo=timezone.utc), 10, 5)
|
|
|
|
async def test_fetch_existing_chunk_keys(self):
|
|
conn = _FakeConn(existing_keys=[("m1@x", 0), ("m1@x", 1)])
|
|
keys = await fetch_existing_chunk_keys(conn, model="bge-m3")
|
|
assert keys == {("m1@x", 0), ("m1@x", 1)}
|
|
|
|
|
|
def _env_row(envelope_id: str, subject: str, ts: datetime = datetime(2025, 8, 1, tzinfo=timezone.utc)):
|
|
entities = [{"type": "headers", "subject": subject, "from": {"name": "A", "address": "a@b.com"}}]
|
|
return {"id": envelope_id, "ts": ts, "raw_ref": f"gmail/2025/08/{envelope_id}.eml", "entities": json.dumps(entities)}
|
|
|
|
|
|
class _FakeConn:
|
|
"""Mirrors documents_ingest's/backfill's `_FakeConn` test style. `existing_keys` were
|
|
"written" under `existing_model` -- a fetch for a different model must not see them,
|
|
same as the real `WHERE model = $1` filter."""
|
|
|
|
def __init__(self, envelopes=None, existing_keys=None, existing_model="bge-m3", execute_results=None):
|
|
self._envelopes = envelopes or []
|
|
self._existing_keys = list(existing_keys or [])
|
|
self._existing_model = existing_model
|
|
self._execute_results = list(execute_results) if execute_results is not None else None
|
|
self.execute_calls: list[tuple] = []
|
|
self.executemany_calls: list[tuple] = []
|
|
self.queries: list[tuple] = []
|
|
|
|
async def fetch(self, query, *params):
|
|
self.queries.append((query, params))
|
|
if "FROM document_chunk" in query:
|
|
if params and params[0] != self._existing_model:
|
|
return []
|
|
return [{"envelope_id": eid, "chunk_index": idx} for eid, idx in self._existing_keys]
|
|
if "FROM envelope" in query:
|
|
return self._envelopes
|
|
raise AssertionError(f"unexpected query: {query}")
|
|
|
|
async def execute(self, query, *params):
|
|
self.execute_calls.append(params)
|
|
if self._execute_results is not None:
|
|
return self._execute_results.pop(0)
|
|
return "INSERT 0 1"
|
|
|
|
async def executemany(self, query, rows):
|
|
self.executemany_calls.append((query, list(rows)))
|
|
|
|
async def close(self):
|
|
pass
|
|
|
|
|
|
class _FakeEmbedResponse:
|
|
def __init__(self, payload, status=200):
|
|
self._payload = payload
|
|
self._status = status
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *exc):
|
|
return False
|
|
|
|
async def json(self):
|
|
return self._payload
|
|
|
|
def raise_for_status(self):
|
|
if self._status >= 400:
|
|
raise aiohttp.ClientConnectionError(f"simulated HTTP {self._status}")
|
|
|
|
|
|
class _FakeTagsResponse:
|
|
def __init__(self, status=200):
|
|
self.status = status
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *exc):
|
|
return False
|
|
|
|
|
|
class _FakeOllamaSession:
|
|
"""`health_up` doubles as "is the backend alive": it answers the /api/tags probe that
|
|
`embed_batch_resilient` uses to tell a dead backend (-> give up, breaker counts it) from a
|
|
poison input on a live one (-> bisect, breaker unaffected). A test that wants the breaker
|
|
to trip must therefore set health_up=False, not merely fail the batches."""
|
|
|
|
def __init__(self, dim=1024, fail_batches=False, health_up=True, fail_pattern=None,
|
|
fail_texts=None):
|
|
self._dim = dim
|
|
self._fail_batches = fail_batches
|
|
self._health_up = health_up
|
|
# Per-request failure sequence (True = this batch fails); exhausting it falls back to
|
|
# `fail_batches`. Lets a test interleave failures and successes to exercise the breaker's
|
|
# "consecutive" semantics rather than a plain total.
|
|
self._fail_pattern = list(fail_pattern) if fail_pattern is not None else None
|
|
# Substrings that poison whatever batch they land in — the all-or-nothing property of
|
|
# /api/embed that bisection exists to work around.
|
|
self._fail_texts = list(fail_texts or ())
|
|
self.requests: list[dict] = []
|
|
self.health_checks = 0
|
|
self.closed = False
|
|
|
|
def _should_fail(self, texts) -> bool:
|
|
if any(needle in t for t in texts for needle in self._fail_texts):
|
|
return True
|
|
if self._fail_pattern:
|
|
return self._fail_pattern.pop(0)
|
|
return self._fail_batches
|
|
|
|
def post(self, url, json, timeout=None):
|
|
assert url.endswith("/api/embed")
|
|
self.requests.append({"url": url, "json": json})
|
|
if self._should_fail(json["input"]):
|
|
return _FakeEmbedResponse({}, status=500)
|
|
n = len(json["input"])
|
|
return _FakeEmbedResponse({"embeddings": [[0.01] * self._dim for _ in range(n)]})
|
|
|
|
def get(self, url, timeout=None):
|
|
self.health_checks += 1
|
|
return _FakeTagsResponse(200 if self._health_up else 500)
|
|
|
|
async def close(self):
|
|
self.closed = True
|
|
|
|
|
|
def _write_eml(archive_root, rel_path: str, raw: bytes) -> None:
|
|
dest = archive_root / rel_path
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
dest.write_bytes(raw)
|
|
|
|
|
|
def _patch_run(monkeypatch, conn, ollama_session) -> None:
|
|
async def _fake_connect(dsn):
|
|
return conn
|
|
monkeypatch.setattr("mail_body_ingest.ingest.asyncpg.connect", _fake_connect)
|
|
|
|
def _fake_session_factory(*args, **kwargs):
|
|
return ollama_session
|
|
monkeypatch.setattr("mail_body_ingest.ingest.aiohttp.ClientSession", _fake_session_factory)
|
|
|
|
|
|
class TestRun:
|
|
def _patch(self, monkeypatch, conn, ollama_session):
|
|
_patch_run(monkeypatch, conn, ollama_session)
|
|
|
|
async def test_dry_run_counts_without_calling_ollama_or_db(self, tmp_path, monkeypatch):
|
|
_write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there"
|
|
))
|
|
conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")])
|
|
ollama = _FakeOllamaSession()
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=False)
|
|
|
|
assert stats["mails_scanned"] == 1
|
|
assert stats["mails_chunked"] == 1
|
|
assert stats["chunks_total"] == 1
|
|
assert stats["chunks_inserted"] == 1
|
|
assert stats["threading_updated"] == 1 # would-update, not applied
|
|
assert ollama.requests == []
|
|
assert conn.execute_calls == []
|
|
assert conn.executemany_calls == []
|
|
|
|
async def test_apply_embeds_inserts_and_updates_threading(self, tmp_path, monkeypatch):
|
|
_write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there"
|
|
))
|
|
conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True)
|
|
|
|
assert stats["chunks_inserted"] == 1
|
|
assert stats["embed_calls"] == 1
|
|
assert len(conn.execute_calls) == 1
|
|
assert len(ollama.requests) == 1
|
|
assert len(conn.executemany_calls) == 1 # threading update batch flushed
|
|
|
|
async def test_idempotent_skips_already_inserted_chunks(self, tmp_path, monkeypatch):
|
|
_write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there"
|
|
))
|
|
conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")], existing_keys=[("m1@x", 0)])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True)
|
|
|
|
assert stats["chunks_already_embedded"] == 1
|
|
assert stats["chunks_inserted"] == 0
|
|
assert ollama.requests == []
|
|
|
|
async def test_rerun_after_apply_inserts_nothing_new(self, tmp_path, monkeypatch):
|
|
_write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there"
|
|
))
|
|
conn1 = _FakeConn(envelopes=[_env_row("m1@x", "hi")])
|
|
ollama1 = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn1, ollama1)
|
|
first = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True)
|
|
assert first["chunks_inserted"] == 1
|
|
|
|
conn2 = _FakeConn(envelopes=[_env_row("m1@x", "hi")], existing_keys=[("m1@x", 0)])
|
|
ollama2 = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn2, ollama2)
|
|
second = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True)
|
|
|
|
assert second["chunks_inserted"] == 0
|
|
assert second["chunks_already_embedded"] == 1
|
|
assert ollama2.requests == []
|
|
|
|
async def test_newsletter_chunk_flagged_and_not_embedded(self, tmp_path, monkeypatch):
|
|
_write_eml(tmp_path, "gmail/2025/08/nl@x.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": "Promo", "List-Unsubscribe": "<mailto:x@y.com>",
|
|
"Date": "Fri, 01 Aug 2025 10:00:00 +0000"},
|
|
"Kup teraz z rabatem!",
|
|
))
|
|
conn = _FakeConn(envelopes=[_env_row("nl@x", "Promo")])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True)
|
|
|
|
assert stats["chunks_newsletter_flagged"] == 1
|
|
assert stats["chunks_inserted"] == 0
|
|
assert ollama.requests == [] # no Ollama call for newsletter chunks
|
|
assert conn.execute_calls[0][3] is None # embedding column is NULL
|
|
assert conn.execute_calls[0][5] == "newsletter"
|
|
|
|
async def test_body_empty_after_quote_strip_still_updates_threading(self, tmp_path, monkeypatch):
|
|
_write_eml(tmp_path, "gmail/2025/08/empty@x.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": "Re: x", "In-Reply-To": "<orig@x.com>",
|
|
"Date": "Fri, 01 Aug 2025 10:00:00 +0000"},
|
|
"On Tue wrote:\n> everything is quoted",
|
|
))
|
|
conn = _FakeConn(envelopes=[_env_row("empty@x", "Re: x")])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True)
|
|
|
|
assert stats["body_empty"] == 1
|
|
assert stats["chunks_total"] == 0
|
|
assert stats["threading_updated"] == 1
|
|
assert len(conn.executemany_calls) == 1
|
|
patch = json.loads(conn.executemany_calls[0][1][0][1])
|
|
assert patch[0]["in_reply_to"] == "orig@x.com"
|
|
|
|
async def test_missing_file_counted(self, tmp_path, monkeypatch):
|
|
conn = _FakeConn(envelopes=[_env_row("missing@x", "hi")])
|
|
ollama = _FakeOllamaSession()
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=False)
|
|
|
|
assert stats["missing_file"] == 1
|
|
assert stats["mails_scanned"] == 1
|
|
|
|
async def test_ollama_offline_batch_error_isolated_not_aborting(self, tmp_path, monkeypatch):
|
|
_write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there"
|
|
))
|
|
conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")])
|
|
ollama = _FakeOllamaSession(dim=1024, fail_batches=True, health_up=False)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True,
|
|
embed_retries=0)
|
|
|
|
assert stats["chunks_errors"] == 1
|
|
assert stats["chunks_inserted"] == 0
|
|
assert stats["chunks_total"] == (
|
|
stats["chunks_inserted"] + stats["chunks_newsletter_flagged"] + stats["chunks_already_embedded"]
|
|
+ stats["chunks_conflict_skipped"] + stats["chunks_errors"]
|
|
)
|
|
|
|
def _write_n_mails(self, tmp_path, n: int) -> list:
|
|
rows = []
|
|
for i in range(n):
|
|
eid = f"m{i}@x"
|
|
_write_eml(tmp_path, f"gmail/2025/08/{eid}.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": f"hi {i}", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"},
|
|
f"hello there number {i}",
|
|
))
|
|
rows.append(_env_row(eid, f"hi {i}"))
|
|
return rows
|
|
|
|
async def test_breaker_aborts_after_consecutive_embed_failures(self, tmp_path, monkeypatch):
|
|
"""A dead backend must stop the run, not let it parse the rest of the archive into
|
|
chunks_errors (plan §9 Etap B: 200k+ mails behind a single-threaded parse)."""
|
|
conn = _FakeConn(envelopes=self._write_n_mails(tmp_path, 8))
|
|
ollama = _FakeOllamaSession(dim=1024, fail_batches=True, health_up=False)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
with pytest.raises(EmbedBackendUnavailableError):
|
|
await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True,
|
|
batch_size=1, max_embed_failures=5, embed_retries=0)
|
|
|
|
# Stopped at the 5th failed batch — mails 6-8 were never parsed, let alone embedded.
|
|
assert len(ollama.requests) == 5
|
|
assert conn.execute_calls == [] # nothing inserted: failed chunks stay retryable
|
|
|
|
async def test_breaker_counter_resets_on_successful_batch(self, tmp_path, monkeypatch):
|
|
""""Consecutive", not "total" — flaky batches interleaved with successes must not trip it."""
|
|
conn = _FakeConn(envelopes=self._write_n_mails(tmp_path, 5))
|
|
ollama = _FakeOllamaSession(dim=1024, health_up=False,
|
|
fail_pattern=[True, True, False, True, True])
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True,
|
|
batch_size=1, max_embed_failures=3, embed_retries=0)
|
|
|
|
assert len(ollama.requests) == 5 # ran to completion
|
|
assert stats["chunks_errors"] == 4
|
|
assert stats["chunks_inserted"] == 1
|
|
assert stats["embed_batch_failures"] == 4
|
|
assert stats["chunks_total"] == (
|
|
stats["chunks_inserted"] + stats["chunks_newsletter_flagged"] + stats["chunks_already_embedded"]
|
|
+ stats["chunks_conflict_skipped"] + stats["chunks_errors"]
|
|
)
|
|
|
|
async def test_breaker_disabled_with_zero(self, tmp_path, monkeypatch):
|
|
conn = _FakeConn(envelopes=self._write_n_mails(tmp_path, 6))
|
|
ollama = _FakeOllamaSession(dim=1024, fail_batches=True, health_up=False)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True,
|
|
batch_size=1, max_embed_failures=0, embed_retries=0)
|
|
|
|
assert len(ollama.requests) == 6
|
|
assert stats["chunks_errors"] == 6
|
|
|
|
async def test_breaker_abort_flushes_pending_threading(self, tmp_path, monkeypatch):
|
|
"""Threading appends are Ollama-independent and idempotent — the abort keeps them rather
|
|
than making the next run re-derive them from the same 27 GB read."""
|
|
conn = _FakeConn(envelopes=self._write_n_mails(tmp_path, 8))
|
|
ollama = _FakeOllamaSession(dim=1024, fail_batches=True, health_up=False)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
with pytest.raises(EmbedBackendUnavailableError):
|
|
await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True,
|
|
batch_size=1, max_embed_failures=5, embed_retries=0)
|
|
|
|
assert len(conn.executemany_calls) == 1
|
|
# The 5 mails processed before the breaker tripped, none of the 3 after it.
|
|
assert len(conn.executemany_calls[0][1]) == 5
|
|
|
|
async def test_partial_batch_failure_does_not_advance_the_breaker(self, tmp_path, monkeypatch):
|
|
"""A live backend that chokes on one chunk must not abort a 50k slice. The bad chunk is
|
|
isolated, everything else embeds, and the run continues past max_embed_failures batches."""
|
|
rows = self._write_n_mails(tmp_path, 6)
|
|
# Mail 2's body poisons any batch it lands in; the backend stays healthy throughout.
|
|
ollama = _FakeOllamaSession(dim=1024, health_up=True, fail_texts=["number 2"])
|
|
conn = _FakeConn(envelopes=rows)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True,
|
|
batch_size=1, max_embed_failures=1, embed_retries=0)
|
|
|
|
assert stats["chunks_errors"] == 1 # only the poison chunk
|
|
assert stats["chunks_inserted"] == 5 # the other five got through
|
|
assert stats["embed_items_failed"] == 1
|
|
assert stats["mails_scanned"] == 6 # ran to completion, never aborted
|
|
assert stats["chunks_total"] == (
|
|
stats["chunks_inserted"] + stats["chunks_newsletter_flagged"]
|
|
+ stats["chunks_already_embedded"] + stats["chunks_conflict_skipped"]
|
|
+ stats["chunks_errors"]
|
|
)
|
|
|
|
async def test_poison_chunk_isolated_from_a_full_batch(self, tmp_path, monkeypatch):
|
|
"""Batched: one bad chunk used to cost the whole batch. Bisection keeps the rest."""
|
|
conn = _FakeConn(envelopes=self._write_n_mails(tmp_path, 8))
|
|
ollama = _FakeOllamaSession(dim=1024, health_up=True, fail_texts=["number 5"])
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True,
|
|
batch_size=8, embed_retries=0)
|
|
|
|
assert stats["chunks_inserted"] == 7
|
|
assert stats["chunks_errors"] == 1
|
|
assert stats["embed_calls"] == 1 # one logical batch...
|
|
assert stats["embed_requests_total"] > 1 # ...several HTTP requests to isolate it
|
|
|
|
async def test_embed_timeout_degrades_instead_of_crashing_the_run(self, tmp_path, monkeypatch):
|
|
"""Regression: a bare builtins.TimeoutError from an exhausted ClientTimeout is not an
|
|
aiohttp.ClientError, so it used to escape the handler and kill the run — no breaker, no
|
|
threading flush. Ollama@SOLARIA hangs rather than refusing, so this was reachable."""
|
|
conn = _FakeConn(envelopes=self._write_n_mails(tmp_path, 3))
|
|
|
|
class _HangingSession(_FakeOllamaSession):
|
|
def post(self, url, json, timeout=None):
|
|
self.requests.append({"url": url, "json": json})
|
|
raise TimeoutError("simulated Ollama hang")
|
|
|
|
ollama = _HangingSession(dim=1024, health_up=False)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
with pytest.raises(EmbedBackendUnavailableError):
|
|
await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True,
|
|
batch_size=1, max_embed_failures=2, embed_retries=0)
|
|
|
|
# Reached the breaker (a clean stop) rather than propagating TimeoutError, and the
|
|
# threading work earned before the stop was still committed.
|
|
assert len(conn.executemany_calls) == 1
|
|
|
|
async def test_env_overrides_batch_size(self, monkeypatch):
|
|
monkeypatch.setenv("MAIL_INGEST_BATCH_SIZE", "16")
|
|
assert _env_num("MAIL_INGEST_BATCH_SIZE", 64, int) == 16
|
|
|
|
async def test_malformed_env_is_a_loud_failure(self, monkeypatch):
|
|
monkeypatch.setenv("MAIL_INGEST_BATCH_SIZE", "sixty-four")
|
|
with pytest.raises(SystemExit):
|
|
_env_num("MAIL_INGEST_BATCH_SIZE", 64, int)
|
|
|
|
async def test_embed_ms_per_chunk_is_per_chunk_not_per_batch(self):
|
|
stats = _new_stats()
|
|
stats["embed_seconds_total"] = 2.0
|
|
stats["embed_texts_ok"] = 100
|
|
assert _embed_ms_per_chunk(stats) == 20.0
|
|
|
|
async def test_embed_ms_per_chunk_handles_zero_embeds(self):
|
|
assert _embed_ms_per_chunk(_new_stats()) == 0.0
|
|
|
|
async def test_dimension_mismatch_aborts(self, tmp_path, monkeypatch):
|
|
_write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there"
|
|
))
|
|
conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")])
|
|
ollama = _FakeOllamaSession(dim=768) # wrong dim
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
with pytest.raises(EmbeddingDimensionError):
|
|
await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True)
|
|
|
|
async def test_conflict_skipped_counted_separately(self, tmp_path, monkeypatch):
|
|
_write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there"
|
|
))
|
|
conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")], execute_results=["INSERT 0 0"])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True)
|
|
|
|
assert stats["chunks_conflict_skipped"] == 1
|
|
assert stats["chunks_inserted"] == 0
|
|
|
|
async def test_threading_already_present_is_skipped(self, tmp_path, monkeypatch):
|
|
_write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"}, "hello there"
|
|
))
|
|
row = _env_row("m1@x", "hi")
|
|
entities = json.loads(row["entities"]) + [{"type": "threading", "in_reply_to": None, "references": []}]
|
|
row["entities"] = json.dumps(entities)
|
|
conn = _FakeConn(envelopes=[row])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
self._patch(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True)
|
|
|
|
assert stats["threading_already_present"] == 1
|
|
assert stats["threading_updated"] == 0
|
|
assert conn.executemany_calls == []
|
|
|
|
|
|
class TestNulBytes:
|
|
"""2007-era mail in the Etap B run (slice offset 50000) carried stray NULs; postgres answered
|
|
with CharacterNotInRepertoireError ('invalid byte sequence for encoding UTF8: 0x00') and the
|
|
chunks were lost. The strip happens before chunking, so the embedded text and the stored text
|
|
are the same string."""
|
|
|
|
async def test_body_with_nul_embeds_and_inserts_clean_text(self, tmp_path, monkeypatch):
|
|
_write_eml(tmp_path, "gmail/2007/03/nul@x.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": "stary mail", "Date": "Thu, 01 Mar 2007 10:00:00 +0000"},
|
|
"przed\x00 i\x00 po",
|
|
))
|
|
row = _env_row("nul@x", "stary mail", ts=datetime(2007, 3, 1, tzinfo=timezone.utc))
|
|
row["raw_ref"] = "gmail/2007/03/nul@x.eml"
|
|
conn = _FakeConn(envelopes=[row])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
_patch_run(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True)
|
|
|
|
assert stats["chunks_inserted"] == 1
|
|
assert stats["chunks_errors"] == 0
|
|
assert stats["nul_bytes_stripped"] == 2
|
|
assert stats["mails_nul_sanitized"] == 1
|
|
|
|
embedded = ollama.requests[0]["json"]["input"]
|
|
inserted_text = conn.execute_calls[0][2]
|
|
assert "\x00" not in inserted_text
|
|
assert [t for t in embedded if "\x00" in t] == []
|
|
# Same string embedded as stored -- that is the whole point of stripping this early.
|
|
assert embedded == [inserted_text]
|
|
assert "przed i po" in inserted_text
|
|
|
|
async def test_ordinary_body_is_untouched_and_counters_stay_zero(self, tmp_path, monkeypatch):
|
|
_write_eml(tmp_path, "gmail/2025/08/m1@x.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": "hi", "Date": "Fri, 01 Aug 2025 10:00:00 +0000"},
|
|
"Zwykły tekst\n\nz akapitem i tabulatorem\ti końcem.",
|
|
))
|
|
conn = _FakeConn(envelopes=[_env_row("m1@x", "hi")])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
_patch_run(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True)
|
|
|
|
assert stats["nul_bytes_stripped"] == 0
|
|
assert stats["mails_nul_sanitized"] == 0
|
|
assert stats["chunks_inserted"] == 1
|
|
inserted_text = conn.execute_calls[0][2]
|
|
assert "Zwykły tekst\n\nz akapitem i tabulatorem\ti końcem." in inserted_text
|
|
|
|
async def test_body_of_only_nuls_counts_as_empty_not_as_a_chunk(self, tmp_path, monkeypatch):
|
|
_write_eml(tmp_path, "gmail/2007/03/nul@x.eml", _plain_eml(
|
|
{"From": "a@b.com", "Subject": "pusty", "Date": "Thu, 01 Mar 2007 10:00:00 +0000"},
|
|
"\x00\x00\x00",
|
|
))
|
|
row = _env_row("nul@x", "pusty", ts=datetime(2007, 3, 1, tzinfo=timezone.utc))
|
|
row["raw_ref"] = "gmail/2007/03/nul@x.eml"
|
|
conn = _FakeConn(envelopes=[row])
|
|
ollama = _FakeOllamaSession(dim=1024)
|
|
_patch_run(monkeypatch, conn, ollama)
|
|
|
|
stats = await run(dsn="postgresql://fake", archive_root=tmp_path, apply=True)
|
|
|
|
assert stats["nul_bytes_stripped"] == 3
|
|
assert stats["body_empty"] == 1
|
|
assert stats["mails_chunked"] == 0
|
|
assert stats["chunks_total"] == 0
|
|
assert ollama.requests == []
|
|
|
|
def test_threading_ids_are_nul_stripped_before_jsonb(self):
|
|
# Same crash class, different column: json.dumps turns a NUL into \\u0000, which postgres
|
|
# jsonb rejects too. extract_threading normalizes before the id tokens are cut out.
|
|
msg = parse_message(_plain_eml(
|
|
{"From": "a@b.com", "In-Reply-To": "<msg\x001@x.com>",
|
|
"References": "<msg\x000@x.com>"}, "hi",
|
|
))
|
|
threading = extract_threading(msg)
|
|
assert threading["in_reply_to"] == "msg1@x.com"
|
|
assert threading["references"] == ["msg0@x.com"]
|