FabledScribe/tests/test_integration_embedding_yields_to_delete.py

"""#3262 against a real Postgres: the embedder loses the race, it doesn't run it.

The reported failure was a deadlock — `DELETE FROM rulebooks` killed by the
server while a detached `upsert_rule_embedding` held the other half of the
cycle. It cannot be reproduced with mocks, because there is nothing to
deadlock: the whole bug lives in the ORDER two transactions take two row
locks, which only a lock manager can adjudicate.

So each test here holds a real delete open in one transaction and calls the
embedder in another. What is being pinned is that the embedder RETURNS —
promptly, having written nothing. Before the fix it would sit on the chunk
rows waiting for a delete that is itself waiting on the insert's foreign key,
and the test would hang rather than fail, which is why every call carries a
deadline (rule 156).
"""
import asyncio
from unittest.mock import AsyncMock, patch

import pytest
import pytest_asyncio
from sqlalchemy import delete, select

from scribe.models import async_session
from scribe.models.embedding import NoteEmbedding, RuleEmbedding
from scribe.models.note import Note
from scribe.models.rulebook import Rulebook
from scribe.services import embeddings as emb
from scribe.services import rulebooks as rulebooks_svc
from tests.helpers import ensure_user

pytestmark = [pytest.mark.integration, pytest.mark.usefixtures("_dispose_engine")]

OWNER_USERNAME = "embed_lock_owner"

# Generous, because the assertion is "it did not block indefinitely", not "it
# was fast". A machine under load must not turn this into a flake; a genuinely
# blocked embedder never returns at all, so no honest run comes near this.
YIELD_DEADLINE_SECONDS = 20

# What the embedder would write if it wrongly went ahead. Distinct from the
# text the fixture's own create_* wrote, so the assertion cannot be satisfied
# by rows that were already there.
SENTINEL = "sentinelvector"

ONE_VECTOR = [[0.0] * 384]

# The fixture's own embedding task runs the REAL embedder, which either loads a
# model or gives up; both are bounded well inside this.
SETTLE_DEADLINE_SECONDS = 30


@pytest_asyncio.fixture
async def seeded():
    """A rule and a note to race against.

    CLEANED AT SETUP, NOT TEARDOWN — the same constraint #3241 hit and the
    reason this file exists. `create_rule` fires its own detached embedding
    task; a teardown that deleted the rulebook would be racing exactly the
    thing under test, on a loop that is closing.
    """
    async with async_session() as s:
        owner = await ensure_user(s, OWNER_USERNAME)
        uid = owner.id
        await s.commit()
        for book in (await s.execute(
            select(Rulebook).where(Rulebook.owner_user_id == uid)
        )).scalars().all():
            await s.delete(book)
        for note in (await s.execute(
            select(Note).where(Note.user_id == uid)
        )).scalars().all():
            await s.delete(note)
        await s.commit()

    book = await rulebooks_svc.create_rulebook(uid, "Lock fixtures")
    topic = await rulebooks_svc.create_topic(book.id, uid, "locks")
    rule = await rulebooks_svc.create_rule(
        topic.id, uid, "A rule with vectors",
        "Something for the embedder to index.",
    )
    async with async_session() as s:
        note = Note(user_id=uid, title="A note with vectors", body="Body text.")
        s.add(note)
        await s.commit()
        note_id = note.id

    await _settle_detached_writes()
    return {"uid": uid, "book_id": book.id, "rule_id": rule.id, "note_id": note_id}


async def _settle_detached_writes() -> None:
    """Let `create_rule`'s own fire-and-forget embedding task finish.

    It is the same detached write these tests are about, aimed at the same
    rule, and left in flight it would land in the middle of an assertion about
    that rule's rows. Bounded, and a timeout is not a failure — the tests below
    carry their own deadlines, and this is only tidying the start line.
    """
    pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
    if pending:
        await asyncio.wait(pending, timeout=SETTLE_DEADLINE_SECONDS)


async def _rule_chunks(rule_id: int) -> list[str]:
    async with async_session() as s:
        return list((await s.execute(
            select(RuleEmbedding.chunk_text).where(RuleEmbedding.rule_id == rule_id)
        )).scalars().all())


async def _note_chunks(note_id: int) -> list[str]:
    async with async_session() as s:
        return list((await s.execute(
            select(NoteEmbedding.chunk_text).where(NoteEmbedding.note_id == note_id)
        )).scalars().all())


async def test_a_rule_refresh_yields_to_a_delete_cascading_from_its_rulebook(seeded):
    """The reported case, exactly: the delete lands on the RULEBOOK and reaches
    the rule through two cascades, which is why nothing on the rule's own write
    path could have seen it coming."""
    async with async_session() as blocker:
        # Uncommitted on purpose — the cascade's locks are held for as long as
        # this transaction stays open, which is the state the embedder must
        # decline to fight over.
        await blocker.execute(delete(Rulebook).where(Rulebook.id == seeded["book_id"]))
        try:
            with patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)):
                await asyncio.wait_for(
                    emb.upsert_rule_embedding(
                        seeded["rule_id"], SENTINEL, f"{SENTINEL} statement",
                    ),
                    timeout=YIELD_DEADLINE_SECONDS,
                )
        finally:
            await blocker.rollback()

    assert not any(SENTINEL in text for text in await _rule_chunks(seeded["rule_id"])), \
        "the embedder wrote into a rule that was being deleted"


async def test_a_note_refresh_yields_to_a_delete_of_the_note(seeded):
    """The note twin, which #3262 recorded as unverified. Notes are soft-deleted
    day to day, so the hard delete a trash purge issues is the one that can put
    a lock on the row while a refresh is in flight."""
    async with async_session() as blocker:
        await blocker.execute(delete(Note).where(Note.id == seeded["note_id"]))
        try:
            with patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)):
                await asyncio.wait_for(
                    emb.upsert_note_embedding(
                        seeded["note_id"], seeded["uid"], SENTINEL, f"{SENTINEL} body",
                    ),
                    timeout=YIELD_DEADLINE_SECONDS,
                )
        finally:
            await blocker.rollback()

    assert not any(SENTINEL in text for text in await _note_chunks(seeded["note_id"])), \
        "the embedder wrote into a note that was being deleted"


async def test_an_uncontended_refresh_still_writes(seeded):
    """The guard against the cheapest possible false pass: a claim that never
    succeeds would satisfy both tests above while quietly ending semantic
    search."""
    with patch.object(emb, "get_embeddings", AsyncMock(return_value=ONE_VECTOR)):
        await emb.upsert_rule_embedding(
            seeded["rule_id"], SENTINEL, f"{SENTINEL} statement",
        )

    assert any(SENTINEL in text for text in await _rule_chunks(seeded["rule_id"])), \
        "an unlocked rule was not embedded"