microsoft--agent-lightning/tests/store/test_collection.py

# Copyright (c) Microsoft. All rights reserved.

from __future__ import annotations

import asyncio
import threading
import time
from typing import (
    TYPE_CHECKING,
    Any,
    Awaitable,
    Callable,
    Dict,
    Iterable,
    List,
    Literal,
    Mapping,
    Sequence,
    Tuple,
    Union,
    cast,
)
from uuid import uuid4

import pydantic
import pytest
from pydantic import BaseModel

import agentlightning.store.collection.memory as memory_module
from agentlightning.store.collection import DequeBasedQueue, DictBasedKeyValue, KeyValue, ListBasedCollection
from agentlightning.store.collection.base import Collection
from agentlightning.store.collection.memory import _item_matches_filters  # pyright: ignore[reportPrivateUsage]
from agentlightning.store.collection.memory import _LoopAwareAsyncLock  # pyright: ignore[reportPrivateUsage]
from agentlightning.store.collection.memory import _ThreadSafeAsyncLock  # pyright: ignore[reportPrivateUsage]
from agentlightning.types import Rollout
from tests.store.conftest import QueueItem, SampleItem, mongo_client_kwargs, mongo_uri

if TYPE_CHECKING:
    from pymongo.asynchronous.database import AsyncDatabase

    from agentlightning.store.collection.mongo import MongoLightningCollections

pytestmark = [pytest.mark.store]


def _build_collection(items: Iterable[SampleItem] = ()) -> ListBasedCollection[SampleItem]:
    return ListBasedCollection(list(items), SampleItem, ("partition", "index"))


BASE_KEY_ORDER: List[Tuple[str, int]] = [
    ("alpha", 1),
    ("alpha", 2),
    ("alpha", 3),
    ("beta", 1),
    ("beta", 2),
    ("gamma", 1),
    ("gamma", 2),
    ("delta", 1),
]


def _key_pairs(items: Sequence[SampleItem]) -> List[Tuple[str, int]]:
    return [(item.partition, item.index) for item in items]


def _sorted_pairs(items: Sequence[SampleItem]) -> List[Tuple[str, int]]:
    return sorted(_key_pairs(items))


def test_list_collection_requires_primary_keys(sample_items: Sequence[SampleItem]) -> None:
    with pytest.raises(ValueError):
        ListBasedCollection(list(sample_items), SampleItem, ())


@pytest.mark.asyncio()
async def test_list_collection_primary_keys(sample_collection: Collection[SampleItem]) -> None:
    assert tuple(sample_collection.primary_keys()) == ("partition", "index")


@pytest.mark.asyncio()
async def test_list_collection_item_type(sample_collection: Collection[SampleItem]) -> None:
    assert sample_collection.item_type() is SampleItem


@pytest.mark.asyncio()
async def test_list_collection_initial_size(
    sample_collection: Collection[SampleItem], sample_items: Sequence[SampleItem]
) -> None:
    assert (await sample_collection.size()) == len(sample_items)


@pytest.mark.asyncio()
async def test_list_collection_repr_contains_model_info(sample_collection: Collection[SampleItem]) -> None:
    result = repr(sample_collection)
    assert sample_collection.__class__.__name__ in result
    assert "SampleItem" in result
    if isinstance(sample_collection, ListBasedCollection):
        assert str(await sample_collection.size()) in result


@pytest.mark.asyncio()
async def test_list_collection_insert_adds_item(sample_collection: Collection[SampleItem]) -> None:
    new_item = SampleItem(partition="omega", index=1, name="omega", status="new")
    await sample_collection.insert([new_item])
    assert (await sample_collection.size()) == 9
    result = await sample_collection.get({"partition": {"exact": "omega"}, "index": {"exact": 1}})
    assert result == new_item


@pytest.mark.asyncio()
async def test_list_collection_insert_duplicate_raises(sample_collection: Collection[SampleItem]) -> None:
    duplicate = SampleItem(partition="alpha", index=1, name="dup", status="new")
    with pytest.raises(ValueError):
        await sample_collection.insert([duplicate])


@pytest.mark.asyncio()
async def test_list_collection_insert_rejects_duplicate_payload(sample_collection: Collection[SampleItem]) -> None:
    """Ensure duplicate items within the same insert batch are rejected."""
    starting_size = await sample_collection.size()
    dup_a = SampleItem(partition="omega", index=1, name="dup-a", status="new")
    dup_b = SampleItem(partition="omega", index=1, name="dup-b", status="new")

    with pytest.raises(ValueError, match=r"(duplicated|Duplicated) primary key"):
        await sample_collection.insert([dup_a, dup_b])

    if isinstance(sample_collection, ListBasedCollection):
        # Only ListBasedCollection supports this rejecting duplicate items within the same insert batch.
        assert await sample_collection.size() == starting_size
        assert await sample_collection.get({"partition": {"exact": "omega"}}) is None


@pytest.mark.asyncio()
async def test_list_collection_insert_wrong_type(sample_collection: Collection[SampleItem]) -> None:
    class Another(BaseModel):
        partition: str
        index: int

    wrong = Another(partition="omega", index=5)
    with pytest.raises(TypeError):
        await sample_collection.insert([wrong])  # type: ignore[arg-type]


@pytest.mark.asyncio()
async def test_list_collection_update_existing(sample_collection: Collection[SampleItem]) -> None:
    updated = SampleItem(partition="alpha", index=1, name="updated", status="new")
    await sample_collection.update([updated])
    result = await sample_collection.get({"partition": {"exact": "alpha"}, "index": {"exact": 1}})
    assert result == updated


@pytest.mark.asyncio()
async def test_list_collection_update_missing_raises(sample_collection: Collection[SampleItem]) -> None:
    missing = SampleItem(partition="omega", index=99, name="missing", status="lost")
    with pytest.raises(ValueError):
        await sample_collection.update([missing])


@pytest.mark.asyncio()
async def test_list_collection_delete_existing(sample_collection: Collection[SampleItem]) -> None:
    target = SampleItem(partition="alpha", index=1, name="ignored", status="new")
    await sample_collection.delete([target])
    assert (await sample_collection.size()) == 7
    result = await sample_collection.get({"partition": {"exact": "alpha"}, "index": {"exact": 1}})
    assert result is None


@pytest.mark.asyncio()
async def test_list_collection_delete_missing_raises(sample_collection: Collection[SampleItem]) -> None:
    missing = SampleItem(partition="omega", index=3, name="x", status="y")
    with pytest.raises(ValueError):
        await sample_collection.delete([missing])


@pytest.mark.asyncio()
async def test_list_collection_upsert_inserts_when_missing(sample_collection: Collection[SampleItem]) -> None:
    created = SampleItem(partition="omega", index=4, name="new", status="queued")
    await sample_collection.upsert([created])
    assert (await sample_collection.size()) == 9
    fetched = await sample_collection.get({"partition": {"exact": "omega"}, "index": {"exact": 4}})
    assert fetched == created


@pytest.mark.asyncio()
async def test_list_collection_upsert_updates_when_existing(sample_collection: Collection[SampleItem]) -> None:
    replacement = SampleItem(partition="beta", index=2, name="replacement", status="done")
    await sample_collection.upsert([replacement])
    assert (await sample_collection.size()) == 8
    fetched = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 2}})
    assert fetched == replacement


@pytest.mark.asyncio()
async def test_list_collection_upsert_get_or_insert_semantics(sample_collection: Collection[SampleItem]) -> None:
    filters: Mapping[str, Any] = {"partition": {"exact": "beta"}, "index": {"exact": 2}}
    original = await sample_collection.get(filters)
    assert original is not None

    replacement = SampleItem(
        partition="beta",
        index=2,
        name="replacement",
        status="queued",
        tags=["patched"],
        score=999,
        rank=999,
        updated_time=99.0,
        payload={"priority": 99},
        metadata="replacement",
    )

    await sample_collection.upsert([replacement], update_fields=[])

    fetched = await sample_collection.get(filters)
    assert fetched == original and fetched is not None
    assert fetched.name == original.name
    assert fetched.status == original.status


@pytest.mark.asyncio()
async def test_list_collection_upsert_updates_selected_fields(sample_collection: Collection[SampleItem]) -> None:
    filters: Mapping[str, Any] = {"partition": {"exact": "beta"}, "index": {"exact": 1}}
    original = await sample_collection.get(filters)
    assert original is not None

    incoming = SampleItem(
        partition="beta",
        index=1,
        name="beta-incoming",
        status="in-progress",
        tags=["different"],
        score=-1.0,
        rank=42,
        updated_time=123.45,
        payload={"priority": -1},
        metadata="incoming",
    )

    await sample_collection.upsert([incoming], update_fields=["status", "updated_time"])

    fetched = await sample_collection.get(filters)
    assert fetched is not None
    assert fetched.status == incoming.status
    assert fetched.updated_time == incoming.updated_time
    # Ensure unspecified fields (e.g. name/tags) remain the same as the original document.
    assert fetched.name == original.name
    assert fetched.tags == original.tags


@pytest.mark.asyncio()
async def test_list_collection_update_returns_mutated_items(sample_collection: Collection[SampleItem]) -> None:
    replacements = [
        SampleItem(partition="alpha", index=1, name="alpha-new", status="patched"),
        SampleItem(partition="delta", index=1, name="delta-new", status="patched"),
    ]

    returned = await sample_collection.update(replacements)
    assert list(returned) == replacements

    for expected in replacements:
        fetched = await sample_collection.get(
            {"partition": {"exact": expected.partition}, "index": {"exact": expected.index}}
        )
        assert fetched == expected

    original_beta = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 1}})
    assert original_beta is not None

    partial_payload = SampleItem(
        partition="beta", index=1, name="ignored", status="partial", metadata="updated-metadata"
    )
    partial_returned = await sample_collection.update([partial_payload], update_fields=["status", "metadata"])
    assert len(partial_returned) == 1

    fetched_partial = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 1}})
    assert fetched_partial == partial_returned[0]
    assert fetched_partial is not None
    assert fetched_partial.status == partial_payload.status
    assert fetched_partial.metadata == partial_payload.metadata
    assert fetched_partial.name == original_beta.name


@pytest.mark.asyncio()
async def test_list_collection_upsert_returns_mutated_items(sample_collection: Collection[SampleItem]) -> None:
    new_item = SampleItem(partition="omega", index=99, name="omega-new", status="queued")
    inserted = await sample_collection.upsert([new_item])
    assert list(inserted) == [new_item]

    fetched_new = await sample_collection.get({"partition": {"exact": "omega"}, "index": {"exact": 99}})
    assert fetched_new == new_item

    original_existing = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 2}})
    assert original_existing is not None

    incoming = SampleItem(
        partition="beta",
        index=2,
        name="beta-incoming-new-name",
        status="processing",
        tags=["beta", "patched"],
    )
    updated = await sample_collection.upsert([incoming], update_fields=["status", "tags"])
    assert len(updated) == 1

    fetched_existing = await sample_collection.get({"partition": {"exact": "beta"}, "index": {"exact": 2}})
    assert fetched_existing == updated[0]
    assert fetched_existing is not None
    assert fetched_existing.status == incoming.status
    assert fetched_existing.tags == incoming.tags
    assert fetched_existing.name == original_existing.name


@pytest.mark.asyncio()
async def test_list_collection_delete_multiple_items(sample_collection: Collection[SampleItem]) -> None:
    await sample_collection.delete(
        [
            SampleItem(partition="alpha", index=1, name="", status=""),
            SampleItem(partition="beta", index=1, name="", status=""),
        ]
    )
    assert (await sample_collection.size()) == 6


@pytest.mark.asyncio()
async def test_list_collection_insert_accepts_tuple_sequence(
    sample_collection: Collection[SampleItem],
) -> None:
    extra = (
        SampleItem(partition="tuple", index=1, name="a", status="pending"),
        SampleItem(partition="tuple", index=2, name="b", status="pending"),
    )
    await sample_collection.insert(extra)
    assert (await sample_collection.size()) == 10
    fetched = await sample_collection.query(filter={"partition": {"exact": "tuple"}})
    assert _sorted_pairs(fetched.items) == [("tuple", 1), ("tuple", 2)]


@pytest.mark.asyncio()
async def test_list_collection_query_without_filters_returns_all(
    sample_collection: Collection[SampleItem],
) -> None:
    result = await sample_collection.query()
    assert result.total == 8
    assert len(result.items) == 8


@pytest.mark.asyncio()
@pytest.mark.parametrize(
    ("filters", "expected"),
    [
        pytest.param({"status": {"exact": "new"}}, [("alpha", 1), ("beta", 1)], id="exact-single-field"),
        pytest.param(
            {"partition": {"exact": "alpha"}, "index": {"exact": 2}},
            [("alpha", 2)],
            id="exact-multiple-fields",
        ),
        pytest.param(
            {"status": {"within": {"running", "blocked"}}},
            [("alpha", 2), ("alpha", 3), ("gamma", 1)],
            id="within-set",
        ),
        pytest.param(
            {"partition": {"within": ["gamma", "delta"]}},
            [("gamma", 1), ("gamma", 2), ("delta", 1)],
            id="within-list",
        ),
        pytest.param(
            {"name": {"contains": "phase"}},
            [("alpha", 1), ("alpha", 2), ("alpha", 3), ("gamma", 1)],
            id="contains-substring",
        ),
        pytest.param(
            {"tags": {"contains": "urgent"}},
            [("alpha", 1), ("beta", 1)],
            id="contains-list",
        ),
        pytest.param({"metadata": {"contains": "phase"}}, [("alpha", 3), ("gamma", 1)], id="contains-with-none-values"),
        pytest.param({"tags": {"contains": "missing"}}, [], id="contains-no-match"),
        pytest.param({"partition": {"exact": "delta"}}, [("delta", 1)], id="single-exact-match"),
        pytest.param({"missing": {"exact": "value"}}, [], id="exact-missing-field"),
        pytest.param({"score": {"contains": "phase"}}, [], id="contains-typeerror"),
        pytest.param({"name": {"contains": None}}, list(BASE_KEY_ORDER), id="contains-null-check"),
        pytest.param({"status": {"exact": None}}, list(BASE_KEY_ORDER), id="exact-null-no-filter"),
        pytest.param({"status": {"within": 1}}, [], id="within-non-iterable"),
    ],
)
async def test_list_collection_query_filters(
    sample_collection: Collection[SampleItem],
    filters: Dict[str, Dict[str, object]],
    expected: Sequence[Tuple[str, int]],
    request: pytest.FixtureRequest,
) -> None:
    # Mongo implementation raises ValueError for non-iterable values in within filter
    if request.node.callspec.id == "mongo-within-non-iterable":  # type: ignore
        with pytest.raises(ValueError):
            await sample_collection.query(filter=filters)  # type: ignore[arg-type]
        return

    result = await sample_collection.query(filter=filters)  # type: ignore[arg-type]
    assert _sorted_pairs(result.items) == sorted(expected)
    assert result.total == len(expected)


@pytest.mark.asyncio()
@pytest.mark.parametrize(
    ("filters", "filter_logic", "expected"),
    [
        (
            {"status": {"exact": "new"}, "tags": {"contains": "beta"}},
            "and",
            [("beta", 1)],
        ),
        (
            {"status": {"exact": "new"}, "tags": {"contains": "beta"}},
            "or",
            [("alpha", 1), ("beta", 1), ("beta", 2)],
        ),
        (
            {"status": {"exact": "done"}, "tags": {"contains": "core"}},
            "and",
            [("gamma", 2)],
        ),
        (
            {"status": {"exact": "done"}, "tags": {"contains": "core"}},
            "or",
            [("alpha", 1), ("alpha", 2), ("beta", 2), ("gamma", 2)],
        ),
    ],
)
async def test_list_collection_filter_logic(
    sample_collection: Collection[SampleItem],
    filters: Dict[str, Dict[str, object]],
    filter_logic: Literal["and", "or"],
    expected: Sequence[Tuple[str, int]],
) -> None:
    filter_payload = dict(filters)
    filter_payload["_aggregate"] = filter_logic  # type: ignore[index]
    result = await sample_collection.query(filter=filter_payload)  # type: ignore[arg-type]
    assert _sorted_pairs(result.items) == sorted(expected)


@pytest.mark.asyncio()
async def test_list_collection_must_filters_respected_with_or(
    sample_collection: Collection[SampleItem],
) -> None:
    filters = {
        "_aggregate": "or",
        "_must": {"partition": {"exact": "alpha"}},
        "status": {"exact": "done"},
        "tags": {"contains": "urgent"},
    }
    result = await sample_collection.query(filter=filters)  # type: ignore[arg-type]
    assert _sorted_pairs(result.items) == [("alpha", 1)]


@pytest.mark.asyncio()
async def test_list_collection_must_filters_accept_sequence(
    sample_collection: Collection[SampleItem],
) -> None:
    filters = {
        "_aggregate": "or",
        "_must": [
            {"partition": {"exact": "beta"}},
            {"index": {"exact": 2}},
        ],
        "status": {"exact": "new"},
        "tags": {"contains": "beta"},
    }
    result = await sample_collection.query(filter=filters)  # type: ignore[arg-type]
    assert _sorted_pairs(result.items) == [("beta", 2)]


@pytest.mark.asyncio()
async def test_list_collection_must_filters_limit_tree_scan_even_with_or(
    sample_collection: Collection[SampleItem],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    if not isinstance(sample_collection, ListBasedCollection):
        pytest.skip("This test is only valid for pure-memory collections")
    seen: List[Tuple[str, int]] = []
    original = _item_matches_filters

    def tracking(
        item: SampleItem,
        filters: object,
        filter_logic: str,
        must_filters: object | None = None,
    ) -> bool:
        seen.append((item.partition, item.index))
        return original(item, filters, filter_logic, must_filters)  # type: ignore[arg-type]

    monkeypatch.setattr(memory_module, "_item_matches_filters", tracking)

    filters = {
        "_aggregate": "or",
        "_must": {"partition": {"exact": "gamma"}},
        "status": {"exact": "done"},
        "tags": {"contains": "urgent"},
    }
    result = await sample_collection.query(filter=filters)  # type: ignore[arg-type]
    assert _sorted_pairs(result.items) == [("gamma", 2)]
    assert set(seen) == {("gamma", 1), ("gamma", 2)}


@pytest.mark.asyncio()
async def test_list_collection_primary_key_prefix_limits_filter_checks(
    sample_items: Sequence[SampleItem],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    collection = _build_collection(sample_items)
    seen: List[Tuple[str, int]] = []
    original = _item_matches_filters

    def tracking(
        item: SampleItem,
        filters: object,
        filter_logic: str,
        must_filters: object | None = None,
    ) -> bool:
        seen.append((item.partition, item.index))
        return original(item, filters, filter_logic, must_filters)  # type: ignore[arg-type]

    monkeypatch.setattr(memory_module, "_item_matches_filters", tracking)

    filters = {"partition": {"exact": "alpha"}, "index": {"within": {1, 2}}}
    result = await collection.query(filter=filters)  # type: ignore[arg-type]
    assert _sorted_pairs(result.items) == [("alpha", 1), ("alpha", 2)]
    assert set(seen) == {("alpha", 1), ("alpha", 2), ("alpha", 3)}


@pytest.mark.asyncio()
async def test_list_collection_full_primary_key_avoids_tree_scan(
    sample_collection: Collection[SampleItem],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    if not isinstance(sample_collection, ListBasedCollection):
        pytest.skip("This test is only valid for pure-memory collections")
    call_count = 0
    original_iter_items = (  # pyright: ignore[reportPrivateUsage,reportUnknownMemberType,reportUnknownVariableType]
        ListBasedCollection._iter_items  # pyright: ignore[reportPrivateUsage,reportUnknownMemberType]
    )

    def tracking(
        self: ListBasedCollection[SampleItem],
        root: Mapping[str, object] | None = None,
        filters: object | None = None,
        must_filters: object | None = None,
        filter_logic: str = "and",
    ) -> Iterable[SampleItem]:
        nonlocal call_count
        call_count += 1
        return original_iter_items(self, root, filters, must_filters, filter_logic)  # type: ignore[arg-type]

    monkeypatch.setattr(ListBasedCollection, "_iter_items", tracking)

    filters = {"partition": {"exact": "beta"}, "index": {"exact": 2}}
    result = await sample_collection.query(filter=filters)  # type: ignore[arg-type]
    assert _sorted_pairs(result.items) == [("beta", 2)]
    assert call_count == 0


@pytest.mark.asyncio()
@pytest.mark.parametrize(
    ("sort_by", "sort_order", "limit", "expected"),
    [
        ("name", "asc", 4, [("beta", 2), ("beta", 1), ("alpha", 3), ("delta", 1)]),
        ("name", "desc", 4, [("alpha", 1), ("alpha", 2), ("gamma", 1), ("gamma", 2)]),
        ("rank", "asc", 4, [("beta", 2), ("beta", 1), ("alpha", 2), ("alpha", 1)]),
        ("rank", "desc", 4, [("delta", 1), ("gamma", 2), ("alpha", 3), ("gamma", 1)]),
        ("score", "asc", 4, [("alpha", 3), ("gamma", 2), ("beta", 2), ("delta", 1)]),
        ("score", "desc", 4, [("alpha", 1), ("gamma", 1), ("beta", 1), ("alpha", 2)]),
        (
            "updated_time",
            "asc",
            4,
            (
                [("alpha", 2), ("gamma", 1), ("beta", 2), ("beta", 1)],
                [("beta", 2), ("beta", 1), ("delta", 1), ("alpha", 1)],
            ),
        ),
        (
            "updated_time",
            "desc",
            4,
            (
                [("gamma", 2), ("alpha", 3), ("alpha", 1), ("delta", 1)],
                [("gamma", 1), ("alpha", 2), ("gamma", 2), ("alpha", 3)],
            ),
        ),
    ],
)
async def test_list_collection_sorting(
    sample_collection: Collection[SampleItem],
    sort_by: str,
    sort_order: str,
    limit: int,
    expected: Union[Sequence[Tuple[str, int]], Tuple[Sequence[Tuple[str, int]], ...]],
) -> None:
    result = await sample_collection.query(sort={"name": sort_by, "order": sort_order}, limit=limit)  # type: ignore[arg-type]
    if isinstance(expected, tuple):
        assert any(_key_pairs(result.items) == list(expected) for expected in expected)
    else:
        assert _key_pairs(result.items) == list(expected)


@pytest.mark.asyncio()
async def test_list_collection_sort_by_missing_field_raises(sample_collection: Collection[SampleItem]) -> None:
    with pytest.raises(ValueError):
        await sample_collection.query(sort={"name": "does_not_exist", "order": "asc"})


@pytest.mark.asyncio()
@pytest.mark.parametrize(
    ("limit", "offset", "expected", "total"),
    [
        (1, 0, [("alpha", 1)], 3),
        (2, 1, [("alpha", 2), ("alpha", 3)], 3),
        (-1, 1, [("alpha", 2), ("alpha", 3)], 3),
        (10, 0, [("alpha", 1), ("alpha", 2), ("alpha", 3)], 3),
        (0, 0, [], 3),
        (1, 10, [], 3),
    ],
)
async def test_list_collection_pagination_without_sort(
    sample_collection: Collection[SampleItem],
    limit: int,
    offset: int,
    expected: Sequence[Tuple[str, int]],
    total: int,
) -> None:
    result = await sample_collection.query(filter={"partition": {"exact": "alpha"}}, limit=limit, offset=offset)
    assert _key_pairs(result.items) == list(expected)
    assert result.total == total


@pytest.mark.asyncio()
async def test_list_collection_pagination_with_sort(sample_collection: Collection[SampleItem]) -> None:
    result = await sample_collection.query(sort={"name": "name", "order": "asc"}, limit=2, offset=3)
    assert _key_pairs(result.items) == [("delta", 1), ("gamma", 2)]
    assert result.total == 8


@pytest.mark.asyncio()
async def test_list_collection_limit_unbounded_with_sort(sample_collection: Collection[SampleItem]) -> None:
    result = await sample_collection.query(sort={"name": "name", "order": "asc"}, limit=-1, offset=6)
    assert _key_pairs(result.items) == [("alpha", 2), ("alpha", 1)]
    assert result.total == 8


@pytest.mark.asyncio()
async def test_list_collection_limit_zero_reports_total(sample_collection: Collection[SampleItem]) -> None:
    result = await sample_collection.query(filter={"status": {"exact": "done"}}, limit=0)
    assert result.items == []
    assert result.total == 2


@pytest.mark.asyncio()
async def test_list_collection_offset_beyond_total_returns_empty(
    sample_collection: Collection[SampleItem],
) -> None:
    result = await sample_collection.query(filter={"status": {"exact": "done"}}, offset=10)
    assert result.items == []
    assert result.total == 2


@pytest.mark.asyncio()
async def test_list_collection_query_reports_total_with_limit(
    sample_collection: Collection[SampleItem],
) -> None:
    result = await sample_collection.query(filter={"partition": {"exact": "alpha"}}, limit=1)
    assert result.total == 3
    assert len(result.items) == 1


@pytest.mark.asyncio()
async def test_list_collection_get_returns_first_match(sample_collection: Collection[SampleItem]) -> None:
    item = await sample_collection.get({"status": {"exact": "new"}})
    assert item is not None
    assert (item.partition, item.index) in [("beta", 1), ("alpha", 1)]


@pytest.mark.asyncio()
async def test_list_collection_get_returns_none(sample_collection: Collection[SampleItem]) -> None:
    result = await sample_collection.get({"partition": {"exact": "missing"}})
    assert result is None


@pytest.mark.asyncio()
async def test_list_collection_get_respects_filter_logic(sample_collection: Collection[SampleItem]) -> None:
    filters = {"status": {"exact": "done"}, "tags": {"contains": "urgent"}, "_aggregate": "or"}
    item = await sample_collection.get(filters)  # type: ignore[arg-type]
    assert item is not None
    assert (item.partition, item.index) in [("gamma", 2), ("alpha", 1)]


@pytest.mark.asyncio()
async def test_list_collection_get_honors_sort_by(sample_collection: Collection[SampleItem]) -> None:
    filters = {"partition": {"exact": "alpha"}}
    item = await sample_collection.get(filters, sort={"name": "rank", "order": "asc"})  # type: ignore[arg-type]
    assert item is not None
    assert (item.partition, item.index) == ("alpha", 2)


@pytest.mark.asyncio()
async def test_list_collection_get_honors_sort_order(sample_collection: Collection[SampleItem]) -> None:
    filters = {"partition": {"exact": "alpha"}}
    item = await sample_collection.get(filters, sort={"name": "rank", "order": "desc"})  # type: ignore[arg-type]
    assert item is not None
    assert (item.partition, item.index) == ("alpha", 3)


@pytest.mark.asyncio()
async def test_list_collection_query_handles_large_dataset() -> None:
    bulk_items = [
        SampleItem(
            partition=f"partition-{i % 5}",
            index=i,
            name=f"bulk-{i}",
            status="bulk",
            score=float(i),
            rank=i,
            updated_time=float(i),
        )
        for i in range(1500)
    ]
    collection = _build_collection(bulk_items)
    result = await collection.query(sort={"name": "index", "order": "asc"}, limit=50, offset=100)
    assert result.total == 1500
    assert len(result.items) == 50
    assert result.items[0].index == 100
    assert result.items[-1].index == 149


@pytest.mark.asyncio()
async def test_list_collection_bulk_delete_and_size() -> None:
    items = [SampleItem(partition="bulk", index=i, name=f"item-{i}", status="bulk") for i in range(40)]
    collection = _build_collection(items)
    await collection.delete(items[:20])
    assert (await collection.size()) == 20
    await collection.delete(items[20:])
    assert (await collection.size()) == 0


@pytest.mark.asyncio()
async def test_list_collection_query_rejects_unknown_operator(
    sample_collection: Collection[SampleItem],
) -> None:
    with pytest.raises(ValueError):
        await sample_collection.query(filter={"status": {"invalid": "x"}})  # type: ignore[arg-type]


@pytest.mark.asyncio()
async def test_list_collection_query_result_type() -> None:
    collection = _build_collection([])
    result = await collection.query(filter=None)
    assert result.items == []
    assert result.offset == 0


@pytest.mark.asyncio()
async def test_deque_queue_initial_size(deque_queue: DequeBasedQueue[QueueItem]) -> None:
    assert (await deque_queue.size()) == 3


def test_deque_queue_item_type(deque_queue: DequeBasedQueue[QueueItem]) -> None:
    assert deque_queue.item_type() is QueueItem


@pytest.mark.asyncio()
async def test_deque_queue_has_detects_members(deque_queue: DequeBasedQueue[QueueItem]) -> None:
    assert await deque_queue.has(QueueItem(idx=1))
    assert not await deque_queue.has(QueueItem(idx=99))


@pytest.mark.asyncio()
async def test_deque_queue_enqueue_appends_items(deque_queue: DequeBasedQueue[QueueItem]) -> None:
    items = [QueueItem(idx=3), QueueItem(idx=4)]
    returned = await deque_queue.enqueue(items)
    assert returned == items
    assert (await deque_queue.size()) == 5


@pytest.mark.asyncio()
async def test_deque_queue_enqueue_rejects_wrong_type(deque_queue: DequeBasedQueue[QueueItem]) -> None:
    class Wrong(BaseModel):
        idx: int

    with pytest.raises(TypeError):
        await deque_queue.enqueue([Wrong(idx=9)])  # type: ignore[arg-type]


@pytest.mark.asyncio()
@pytest.mark.parametrize("limit", [1, 2, 5])
async def test_deque_queue_dequeue_respects_limit(deque_queue: DequeBasedQueue[QueueItem], limit: int) -> None:
    result = await deque_queue.dequeue(limit)
    assert len(result) == min(limit, 3)
    assert (await deque_queue.size()) == 3 - min(limit, 3)


@pytest.mark.asyncio()
async def test_deque_queue_dequeue_zero_returns_empty(deque_queue: DequeBasedQueue[QueueItem]) -> None:
    assert await deque_queue.dequeue(0) == []


@pytest.mark.asyncio()
async def test_deque_queue_dequeue_more_than_available(deque_queue: DequeBasedQueue[QueueItem]) -> None:
    result = await deque_queue.dequeue(10)
    assert len(result) == 3
    assert (await deque_queue.size()) == 0


@pytest.mark.asyncio()
async def test_deque_queue_peek_preserves_items(deque_queue: DequeBasedQueue[QueueItem]) -> None:
    snapshot = await deque_queue.peek(2)
    assert [item.idx for item in snapshot] == [0, 1]
    assert (await deque_queue.size()) == 3


@pytest.mark.asyncio()
async def test_deque_queue_peek_zero_returns_empty(deque_queue: DequeBasedQueue[QueueItem]) -> None:
    assert await deque_queue.peek(0) == []


@pytest.mark.asyncio()
async def test_deque_queue_peek_after_partial_dequeue(deque_queue: DequeBasedQueue[QueueItem]) -> None:
    await deque_queue.dequeue(1)
    snapshot = await deque_queue.peek(2)
    assert [item.idx for item in snapshot] == [1, 2]


@pytest.mark.asyncio()
async def test_deque_queue_handles_large_volume() -> None:
    queue = DequeBasedQueue(QueueItem)
    items = [QueueItem(idx=i) for i in range(2000)]
    await queue.enqueue(items)
    assert (await queue.size()) == 2000
    drained = await queue.dequeue(1500)
    assert len(drained) == 1500
    assert (await queue.size()) == 500


@pytest.mark.asyncio()
async def test_dict_key_value_initial_state(dict_key_value: DictBasedKeyValue[str, int]) -> None:
    assert await dict_key_value.size() == 2
    assert await dict_key_value.get("alpha") == 1
    assert await dict_key_value.get("missing") is None


@pytest.mark.asyncio()
async def test_dict_key_value_has_handles_presence(dict_key_value: DictBasedKeyValue[str, int]) -> None:
    assert await dict_key_value.has("alpha")
    assert not await dict_key_value.has("gamma")


@pytest.mark.asyncio()
async def test_dict_key_value_set_updates_and_expands(dict_key_value: DictBasedKeyValue[str, int]) -> None:
    await dict_key_value.set("gamma", 3)
    assert await dict_key_value.size() == 3
    await dict_key_value.set("alpha", 99)
    assert await dict_key_value.get("alpha") == 99
    assert await dict_key_value.size() == 3


@pytest.mark.asyncio()
async def test_dict_key_value_pop_returns_default(dict_key_value: DictBasedKeyValue[str, int]) -> None:
    result = await dict_key_value.pop("beta")
    assert result == 2
    assert await dict_key_value.size() == 1
    result = await dict_key_value.pop("missing", 42)
    assert result == 42
    assert await dict_key_value.size() == 1


@pytest.mark.asyncio()
async def test_key_value_inc_updates_existing(dict_key_value: KeyValue[str, int]) -> None:
    new_value = await dict_key_value.inc("alpha", 2)
    assert new_value == 3
    assert await dict_key_value.get("alpha") == 3


@pytest.mark.asyncio()
async def test_key_value_inc_initializes_missing(dict_key_value: KeyValue[str, int]) -> None:
    new_value = await dict_key_value.inc("gamma", 4)
    assert new_value == 4
    assert await dict_key_value.has("gamma")


@pytest.mark.asyncio()
async def test_key_value_inc_rejects_non_numeric_amount(dict_key_value: KeyValue[str, int]) -> None:
    with pytest.raises(TypeError):
        await dict_key_value.inc("alpha", cast(Any, "invalid"))


@pytest.mark.asyncio()
async def test_key_value_chmax_updates_existing(dict_key_value: KeyValue[str, int]) -> None:
    new_value = await dict_key_value.chmax("alpha", 10)
    assert new_value == 10
    assert await dict_key_value.get("alpha") == 10


@pytest.mark.asyncio()
async def test_key_value_chmax_ignores_smaller(dict_key_value: KeyValue[str, int]) -> None:
    initial = await dict_key_value.get("alpha")
    result = await dict_key_value.chmax("alpha", 0)
    assert result == initial
    assert await dict_key_value.get("alpha") == initial


@pytest.mark.asyncio()
async def test_key_value_chmax_initializes_missing(dict_key_value: KeyValue[str, int]) -> None:
    result = await dict_key_value.chmax("gamma", 7)
    assert result == 7
    assert await dict_key_value.get("gamma") == 7


@pytest.mark.asyncio()
async def test_key_value_chmax_rejects_non_numeric_value(dict_key_value: KeyValue[str, int]) -> None:
    with pytest.raises(TypeError):
        await dict_key_value.chmax("alpha", cast(Any, "wrong"))


@pytest.mark.asyncio()
async def test_dict_key_value_inc_rejects_non_numeric_value(dict_key_value_memory: DictBasedKeyValue[str, Any]) -> None:
    await dict_key_value_memory.set("alpha", cast(Any, "na"))
    with pytest.raises(TypeError):
        await dict_key_value_memory.inc("alpha", 1)


@pytest.mark.asyncio()
async def test_dict_key_value_chmax_rejects_non_numeric_value(
    dict_key_value_memory: DictBasedKeyValue[str, Any],
) -> None:
    await dict_key_value_memory.set("alpha", cast(Any, "na"))
    with pytest.raises(TypeError):
        await dict_key_value_memory.chmax("alpha", 1)


@pytest.mark.mongo
@pytest.mark.asyncio()
async def test_mongo_key_value_inc_rejects_non_numeric_value(temporary_mongo_database: AsyncDatabase[Any]) -> None:
    from agentlightning.store.collection.mongo import MongoBasedKeyValue, MongoClientPool

    async with MongoClientPool[Mapping[str, Any]](
        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
    ) as client_pool:
        key_value = MongoBasedKeyValue[str, int](
            client_pool,
            temporary_mongo_database.name,
            f"kv-inc-{uuid4().hex}",
            "partition-inc",
            str,
            int,
        )
        collection = await key_value.ensure_collection()
        await collection.insert_one(
            {
                "partition_id": "partition-inc",
                "key": "alpha",
                "value": "oops",
            }
        )

        with pytest.raises(TypeError):
            await key_value.inc("alpha", 1)


@pytest.mark.mongo
@pytest.mark.asyncio()
async def test_mongo_key_value_chmax_behaves_like_max(temporary_mongo_database: AsyncDatabase[Any]) -> None:
    from agentlightning.store.collection.mongo import MongoBasedKeyValue, MongoClientPool

    async with MongoClientPool[Mapping[str, Any]](
        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
    ) as client_pool:
        key_value = MongoBasedKeyValue[str, int](
            client_pool,
            temporary_mongo_database.name,
            f"kv-chmax-{uuid4().hex}",
            "partition-chmax",
            str,
            int,
        )
        assert await key_value.chmax("alpha", 5) == 5
        assert await key_value.chmax("alpha", 3) == 5
        assert await key_value.chmax("alpha", 9) == 9
        assert await key_value.get("alpha") == 9


@pytest.mark.mongo
@pytest.mark.asyncio()
async def test_mongo_key_value_chmax_rejects_non_numeric_value(temporary_mongo_database: AsyncDatabase[Any]) -> None:
    from agentlightning.store.collection.mongo import MongoBasedKeyValue, MongoClientPool

    async with MongoClientPool[Mapping[str, Any]](
        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
    ) as client_pool:
        key_value = MongoBasedKeyValue[str, int](
            client_pool,
            temporary_mongo_database.name,
            f"kv-chmax-bad-{uuid4().hex}",
            "partition-chmax-bad",
            str,
            int,
        )
        collection = await key_value.ensure_collection()
        await collection.insert_one(
            {
                "partition_id": "partition-chmax-bad",
                "key": "alpha",
                "value": "oops",
            }
        )

        with pytest.raises((TypeError, pydantic.ValidationError)):
            await key_value.chmax("alpha", 1)


def test_thread_safe_async_lock_blocks_threads() -> None:
    lock = _ThreadSafeAsyncLock()
    allow_second = threading.Event()
    second_has_lock = threading.Event()
    release_first = threading.Event()

    async def first() -> None:
        async with lock:
            allow_second.set()
            release_first.wait()

    async def second() -> None:
        allow_second.wait()
        async with lock:
            second_has_lock.set()

    def thread1() -> None:
        asyncio.run(first())

    def thread2() -> None:
        asyncio.run(second())

    t1 = threading.Thread(target=thread1)
    t2 = threading.Thread(target=thread2)
    t1.start()
    t2.start()

    assert allow_second.wait(timeout=1)
    assert not second_has_lock.wait(0.05)

    release_first.set()
    t1.join(timeout=1)
    t2.join(timeout=1)

    assert second_has_lock.is_set()


@pytest.mark.asyncio()
async def test_thread_safe_async_lock_serializes_async_tasks() -> None:
    lock = _ThreadSafeAsyncLock()
    first_entered = asyncio.Event()
    release_first = asyncio.Event()
    second_acquired = asyncio.Event()

    async def first() -> None:
        async with lock:
            first_entered.set()
            await release_first.wait()

    async def second() -> None:
        await first_entered.wait()
        async with lock:
            second_acquired.set()

    task1 = asyncio.create_task(first())
    task2 = asyncio.create_task(second())

    await first_entered.wait()
    await asyncio.sleep(0)
    assert not second_acquired.is_set()

    release_first.set()
    await asyncio.gather(task1, task2)

    assert second_acquired.is_set()


@pytest.mark.asyncio()
async def test_loop_aware_async_lock_serializes_tasks() -> None:
    lock = _LoopAwareAsyncLock()
    first_entered = asyncio.Event()
    release_first = asyncio.Event()
    second_acquired = asyncio.Event()

    async def first() -> None:
        async with lock:
            first_entered.set()
            await release_first.wait()

    async def second() -> None:
        await first_entered.wait()
        async with lock:
            second_acquired.set()

    task1 = asyncio.create_task(first())
    task2 = asyncio.create_task(second())

    await first_entered.wait()
    await asyncio.sleep(0)
    assert not second_acquired.is_set()

    release_first.set()
    await asyncio.gather(task1, task2)

    assert second_acquired.is_set()


@pytest.mark.asyncio()
async def test_loop_aware_async_lock_reuses_loop_specific_lock() -> None:
    lock = _LoopAwareAsyncLock()
    first_lock: asyncio.Lock | None = None

    async with lock as acquired:
        first_lock = acquired
        assert first_lock.locked()

    assert first_lock is not None and not first_lock.locked()

    async with lock as acquired_again:
        assert acquired_again is first_lock


@pytest.mark.asyncio()
async def test_loop_aware_async_lock_distinguishes_event_loops() -> None:
    lock = _LoopAwareAsyncLock()
    main_loop_lock: asyncio.Lock | None = None

    async with lock as acquired:
        main_loop_lock = acquired

    locks_from_threads: List[asyncio.Lock] = []

    def _worker() -> None:
        async def runner() -> None:
            async with lock as acquired:
                locks_from_threads.append(acquired)

        asyncio.run(runner())

    worker = threading.Thread(target=_worker)
    worker.start()
    worker.join(timeout=2)

    assert worker.is_alive() is False
    assert main_loop_lock is not None
    assert locks_from_threads
    assert locks_from_threads[0] is not main_loop_lock


@pytest.mark.asyncio()
async def test_dict_key_value_does_not_mutate_input_mapping(dict_key_value_data: Dict[str, int]) -> None:
    key_value = DictBasedKeyValue(dict_key_value_data)
    await key_value.set("gamma", 3)  # type: ignore[arg-type]
    await key_value.pop("alpha")  # type: ignore[arg-type]
    assert dict_key_value_data == {"alpha": 1, "beta": 2}


@pytest.mark.asyncio()
async def test_inmemory_atomic_read_only_skips_lock() -> None:
    collections = memory_module.InMemoryLightningCollections(lock_type="asyncio")

    class FailingLock:
        async def __aenter__(self) -> None:
            raise AssertionError("read-only atomic block should not acquire the lock")

        async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
            return None

    collections._lock = {"default": FailingLock()}  # type: ignore[attr-defined]

    async with collections.atomic(mode="r", snapshot=False):
        # Should complete without touching the failing lock.
        assert collections.rollouts is not None


@pytest.mark.asyncio()
async def test_inmemory_atomic_snapshot_or_write_acquires_lock() -> None:
    collections = memory_module.InMemoryLightningCollections(lock_type="asyncio")

    class RecordingLock:
        def __init__(self) -> None:
            self.enter_count = 0
            self.exit_count = 0

        async def __aenter__(self) -> None:
            self.enter_count += 1

        async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
            self.exit_count += 1

    lock = RecordingLock()
    collections._lock = {"default": lock}  # type: ignore[attr-defined]

    async with collections.atomic(mode="rw", snapshot=False):
        assert collections.attempts is not None

    assert (lock.enter_count, lock.exit_count) == (1, 1)

    lock.enter_count = lock.exit_count = 0

    async with collections.atomic(mode="r", snapshot=True):
        assert collections.spans is not None

    assert (lock.enter_count, lock.exit_count) == (1, 1)


@pytest.mark.mongo
@pytest.mark.asyncio()
async def test_mongo_based_sanity_check(temporary_mongo_database: AsyncDatabase[Any]) -> None:
    from agentlightning.store.collection.mongo import (
        MongoBasedCollection,
        MongoBasedKeyValue,
        MongoBasedQueue,
        MongoClientPool,
    )

    async with MongoClientPool[Mapping[str, Any]](
        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
    ) as client_pool:
        collection = MongoBasedCollection[Any](
            client_pool, temporary_mongo_database.name, "test", "test-123", ["rollout_id"], Rollout
        )
        await collection.ensure_collection()

        start_time = time.time()
        await collection.insert(
            [Rollout(rollout_id="test-123", input="test-123", start_time=start_time, status="running")]
        )

        result = await collection.query(filter={"status": {"exact": "running"}})
        assert result.items == [
            Rollout(rollout_id="test-123", input="test-123", start_time=start_time, status="running")
        ]

        rollout_queue = MongoBasedQueue[str](
            client_pool, temporary_mongo_database.name, "rollout_queue", "partition-1", str
        )
        await rollout_queue.ensure_collection()

        await rollout_queue.enqueue(["r1", "r2", "r3"])
        assert await rollout_queue.size() == 3
        assert await rollout_queue.peek(2) == ["r1", "r2"]
        assert await rollout_queue.dequeue(2) == ["r1", "r2"]
        assert await rollout_queue.size() == 1

        span_kv = MongoBasedKeyValue[str, int](
            client_pool, temporary_mongo_database.name, "span_sequence_ids", "partition-1", str, int
        )
        await span_kv.ensure_collection()

        await span_kv.set("span-123", 1)
        assert await span_kv.has("span-123")
        assert await span_kv.get("span-123") == 1
        assert await span_kv.pop("span-123") == 1
        assert not await span_kv.has("span-123")


@pytest.mark.mongo
@pytest.mark.asyncio()
async def test_mongo_ensure_collection_creates_partition_scoped_index(
    temporary_mongo_database: AsyncDatabase[Any],
) -> None:
    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool

    collection_name = f"ensure-{uuid4().hex}"
    async with MongoClientPool[Mapping[str, Any]](
        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
    ) as client_pool:
        collection = MongoBasedCollection[Any](
            client_pool,
            temporary_mongo_database.name,
            collection_name,
            "partition-ensure",
            ["name", "index"],
            SampleItem,
        )
        await collection.ensure_collection()

        unique_index = None
        async for index in await temporary_mongo_database[collection_name].list_indexes():  # type: ignore
            if index["name"] == "uniq_partition_name_index" and index.get("unique"):  # type: ignore
                unique_index = index  # type: ignore
                break

        assert unique_index is not None, "expected unique partition/index key"
        key_pairs = list(unique_index["key"].items())  # type: ignore
        assert key_pairs == [("partition_id", 1), ("name", 1), ("index", 1)]


@pytest.mark.mongo
@pytest.mark.asyncio()
async def test_mongo_ensure_collection_survives_concurrent_calls(temporary_mongo_database: AsyncDatabase[Any]) -> None:
    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool

    collection_name = f"ensure-{uuid4().hex}"

    async def ensure_once() -> None:
        async with MongoClientPool[Mapping[str, Any]](
            mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
        ) as client_pool:
            collection = MongoBasedCollection(
                client_pool,
                temporary_mongo_database.name,
                collection_name,
                "partition-concurrent",
                ["index"],
                SampleItem,
            )
            await collection.ensure_collection()

    await asyncio.gather(*(ensure_once() for _ in range(20)))

    names = await temporary_mongo_database.list_collection_names()
    assert names.count(collection_name) == 1

    unique_indexes = []
    async for index in await temporary_mongo_database[collection_name].list_indexes():  # type: ignore
        if index["name"].startswith("uniq_partition_"):  # type: ignore
            unique_indexes.append(index["name"])  # type: ignore
    assert unique_indexes == ["uniq_partition_index"]


@pytest.mark.mongo
@pytest.mark.asyncio()
async def test_mongo_ensure_collection_repeats_without_altering_indexes(
    temporary_mongo_database: AsyncDatabase[Any],
) -> None:
    from agentlightning.store.collection.mongo import MongoBasedCollection, MongoClientPool

    collection_name = f"ensure-{uuid4().hex}"
    async with MongoClientPool[Mapping[str, Any]](
        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
    ) as client_pool:
        collection = MongoBasedCollection(
            client_pool, temporary_mongo_database.name, collection_name, "partition-repeat", ["index"], SampleItem
        )
        await collection.ensure_collection()
        await collection.ensure_collection()

        unique_indexes = []
        async for index in await temporary_mongo_database[collection_name].list_indexes():  # type: ignore
            if index["name"].startswith("uniq_partition_"):  # type: ignore
                unique_indexes.append((index["name"], list(index["key"].items())))  # type: ignore

        assert unique_indexes == [("uniq_partition_index", [("partition_id", 1), ("index", 1)])]


async def _with_mongo_collections(
    db: AsyncDatabase[Any],
    callback: Callable[[MongoLightningCollections], Awaitable[Any]],
) -> Any:
    from agentlightning.store.collection.mongo import MongoClientPool, MongoLightningCollections

    async with MongoClientPool[Mapping[str, Any]](
        mongo_uri=mongo_uri, mongo_client_kwargs=mongo_client_kwargs
    ) as client_pool:
        collections = MongoLightningCollections(
            client_pool=client_pool,
            database_name=db.name,
            partition_id=f"partition-{uuid4().hex}",
        )
        return await callback(collections)


async def _initialize_counter(collections: MongoLightningCollections, key: str) -> None:
    async def _init(coll: MongoLightningCollections) -> None:
        await coll.span_sequence_ids.set(key, 0)

    await collections.execute(_init, commit=False)


async def _read_counter(collections: MongoLightningCollections, key: str) -> int:
    async def _read(coll: MongoLightningCollections) -> int:
        value = await coll.span_sequence_ids.get(key)
        assert value is not None
        return value

    return await collections.execute(_read, commit=False)


async def _contention_run(
    collections: MongoLightningCollections,
    *,
    key: str,
    commit: bool,
    concurrency: int,
) -> int:
    read_lock = asyncio.Lock()
    ready = asyncio.Event()
    readers_seen = 0

    async def _barrier() -> None:
        nonlocal readers_seen
        async with read_lock:
            readers_seen += 1
            if readers_seen == concurrency:
                ready.set()
        await ready.wait()

    async def worker(_: int) -> None:
        first_attempt = True

        async def callback(coll: MongoLightningCollections) -> None:
            nonlocal first_attempt
            value = await coll.span_sequence_ids.get(key)
            assert value is not None
            if first_attempt:
                await _barrier()
                first_attempt = False
            await asyncio.sleep(0)
            await coll.span_sequence_ids.set(key, value + 1)

        await collections.execute(callback, commit=commit, snapshot=True, mode="rw")

    await asyncio.gather(*(worker(i) for i in range(concurrency)))
    return await _read_counter(collections, key)


@pytest.mark.mongo
@pytest.mark.asyncio()
async def test_mongo_execute_without_commit_allows_lost_updates(
    temporary_mongo_database: AsyncDatabase[Any],
) -> None:
    async def scenario(collections: MongoLightningCollections) -> None:
        counter_key = f"counter-{uuid4().hex}"
        await _initialize_counter(collections, counter_key)
        final_value = await _contention_run(
            collections,
            key=counter_key,
            commit=False,
            concurrency=6,
        )
        assert final_value == 1

    await _with_mongo_collections(temporary_mongo_database, scenario)


@pytest.mark.mongo
@pytest.mark.asyncio()
async def test_mongo_execute_with_commit_retries_until_success(
    temporary_mongo_database: AsyncDatabase[Any],
) -> None:
    async def scenario(collections: MongoLightningCollections) -> None:
        counter_key = f"counter-{uuid4().hex}"
        await _initialize_counter(collections, counter_key)
        final_value = await _contention_run(
            collections,
            key=counter_key,
            commit=True,
            concurrency=6,
        )
        assert final_value == 6

    await _with_mongo_collections(temporary_mongo_database, scenario)