daena/tests/test_persistence.py

from __future__ import annotations

from pathlib import Path

from daena.persistence.lmdb import LMDBBackend
from daena.persistence.models import RecordState, StoredRecord


def test_put_and_get(lmdb_backend: LMDBBackend, sample_record: StoredRecord) -> None:
    lmdb_backend.put(sample_record)
    retrieved = lmdb_backend.get(sample_record.record_id)
    assert retrieved is not None
    assert retrieved.record_id == sample_record.record_id
    assert retrieved.source == "test_source"
    assert retrieved.body == {"message": "hello"}


def test_pending_returns_new_records(
    lmdb_backend: LMDBBackend, sample_record: StoredRecord
) -> None:
    lmdb_backend.put(sample_record)
    pending = lmdb_backend.pending(limit=10)
    assert len(pending) == 1
    assert pending[0].record_id == sample_record.record_id


def test_ack(lmdb_backend: LMDBBackend, sample_record: StoredRecord) -> None:
    lmdb_backend.put(sample_record)
    lmdb_backend.ack(sample_record.record_id)
    retrieved = lmdb_backend.get(sample_record.record_id)
    assert retrieved is not None
    assert retrieved.state == RecordState.delivered
    pending = lmdb_backend.pending()
    assert all(r.record_id != sample_record.record_id for r in pending)


def test_nack(lmdb_backend: LMDBBackend, sample_record: StoredRecord) -> None:
    lmdb_backend.put(sample_record)
    lmdb_backend.nack(sample_record.record_id, error="test error", retry_delay=1.0)
    retrieved = lmdb_backend.get(sample_record.record_id)
    assert retrieved is not None
    assert retrieved.state == RecordState.failed
    assert retrieved.last_error == "test error"
    assert retrieved.next_retry_at is not None


def test_dead(lmdb_backend: LMDBBackend, sample_record: StoredRecord) -> None:
    lmdb_backend.put(sample_record)
    lmdb_backend.dead(sample_record.record_id, error="permanent failure")
    retrieved = lmdb_backend.get(sample_record.record_id)
    assert retrieved is not None
    assert retrieved.state == RecordState.dead


def test_count(lmdb_backend: LMDBBackend, sample_record: StoredRecord) -> None:
    assert lmdb_backend.count() == 0
    lmdb_backend.put(sample_record)
    assert lmdb_backend.count() == 1
    assert lmdb_backend.count(RecordState.pending) == 1


def test_list_by_state(lmdb_backend: LMDBBackend, sample_record: StoredRecord) -> None:
    lmdb_backend.put(sample_record)
    pending = lmdb_backend.list_by_state(RecordState.pending)
    assert len(pending) == 1
    lmdb_backend.ack(sample_record.record_id)
    delivered = lmdb_backend.list_by_state(RecordState.delivered)
    assert len(delivered) == 1


def test_delete_old(lmdb_backend: LMDBBackend, sample_record: StoredRecord) -> None:
    lmdb_backend.put(sample_record)
    lmdb_backend.ack(sample_record.record_id)
    before = "2100-01-01T00:00:00"
    removed = lmdb_backend.delete_old(before, state=RecordState.delivered)
    assert removed == 1
    assert lmdb_backend.get(sample_record.record_id) is None


def test_iter_pending(lmdb_backend: LMDBBackend, sample_record: StoredRecord) -> None:
    lmdb_backend.put(sample_record)
    records = list(lmdb_backend.iter_pending())
    assert len(records) == 1
    assert records[0].record_id == sample_record.record_id


def test_update_state(lmdb_backend: LMDBBackend, sample_record: StoredRecord) -> None:
    lmdb_backend.put(sample_record)
    lmdb_backend.update_state(sample_record.record_id, RecordState.failed, last_error="oops")
    retrieved = lmdb_backend.get(sample_record.record_id)
    assert retrieved is not None
    assert retrieved.state == RecordState.failed
    assert retrieved.last_error == "oops"


def test_stored_record_to_from_dict(sample_record: StoredRecord) -> None:
    d = sample_record.to_dict()
    restored = StoredRecord.from_dict(d)
    assert restored.record_id == sample_record.record_id
    assert restored.source == sample_record.source
    assert restored.body == sample_record.body
    assert restored.state == sample_record.state


def test_ready_for_retry(lmdb_backend: LMDBBackend, sample_record: StoredRecord) -> None:
    lmdb_backend.put(sample_record)
    lmdb_backend.nack(sample_record.record_id, "retry later", retry_delay=0)
    ready = lmdb_backend.ready_for_retry(limit=10)
    assert len(ready) == 1


def test_index_keys_do_not_accumulate(
    lmdb_backend: LMDBBackend, sample_record: StoredRecord
) -> None:
    lmdb_backend.put(sample_record)
    assert lmdb_backend.index_count() == 1
    lmdb_backend.nack(sample_record.record_id, "later", retry_delay=0)
    assert lmdb_backend.index_count() == 2
    lmdb_backend.ack(sample_record.record_id)
    assert lmdb_backend.ready_for_retry() == []
    assert lmdb_backend.pending() == []
    assert lmdb_backend.count(RecordState.failed) == 0
    assert lmdb_backend.count(RecordState.delivered) == 1
    assert lmdb_backend.index_count() == 1


def test_claim_and_recover(lmdb_backend: LMDBBackend, sample_record: StoredRecord) -> None:
    lmdb_backend.put(sample_record)
    claimed = lmdb_backend.claim("out", limit=10)
    assert len(claimed) == 1
    assert claimed[0].deliveries["out"].state == RecordState.delivering
    assert lmdb_backend.pending() == []
    assert lmdb_backend.recover_inflight() == 1
    pending = lmdb_backend.pending()
    assert len(pending) == 1
    assert pending[0].deliveries["out"].state == RecordState.pending


def test_per_sink_ack_leaves_other_sink_pending(lmdb_backend: LMDBBackend) -> None:
    record = StoredRecord(source="s", record_type="t", body={"n": 1}, sinks=["a", "b"])
    lmdb_backend.put(record)
    lmdb_backend.ack(record.record_id, sink="a")
    stored = lmdb_backend.get(record.record_id)
    assert stored is not None
    assert stored.deliveries["a"].state == RecordState.delivered
    assert stored.deliveries["b"].state == RecordState.pending
    assert lmdb_backend.count(RecordState.pending) == 1
    assert lmdb_backend.count(RecordState.delivered) == 1
    assert lmdb_backend.pending(sink="a") == []
    assert len(lmdb_backend.pending(sink="b")) == 1
    removed = lmdb_backend.delete_old("2100-01-01T00:00:00+00:00", state=RecordState.delivered)
    assert removed == 0
    assert lmdb_backend.get(record.record_id) is not None


def test_replay_dead_becomes_pending(
    lmdb_backend: LMDBBackend, sample_record: StoredRecord
) -> None:
    lmdb_backend.put(sample_record)
    lmdb_backend.dead(sample_record.record_id, error="nope")
    replayed = lmdb_backend.replay(record_id=sample_record.record_id)
    stored = lmdb_backend.get(sample_record.record_id)
    assert replayed == 1
    assert stored is not None
    assert stored.deliveries["out"].state == RecordState.pending
    assert stored.deliveries["out"].retry_count == 0
    assert stored.deliveries["out"].last_error is None


def test_delete_expired_keeps_inflight(lmdb_backend: LMDBBackend) -> None:
    record = StoredRecord(source="s", record_type="t", body={}, sinks=["a", "b"])
    lmdb_backend.put(record)
    lmdb_backend.ack(record.record_id, sink="a")
    lmdb_backend.dead(record.record_id, "boom", sink="b")
    removed = lmdb_backend.delete_expired("2100-01-01T00:00:00+00:00", "2000-01-01T00:00:00+00:00")
    assert removed == 0
    removed = lmdb_backend.delete_expired("2100-01-01T00:00:00+00:00", "2100-01-01T00:00:00+00:00")
    assert removed == 1
    assert lmdb_backend.get(record.record_id) is None


def test_migrates_legacy_record(tmp_db_path: Path) -> None:
    import json

    import lmdb

    tmp_db_path.mkdir(parents=True)
    env = lmdb.open(str(tmp_db_path), map_size=10_485_760, subdir=True)
    payload = {
        "record_id": "legacy",
        "source": "s",
        "record_type": "t",
        "body": {"message": "old"},
        "metadata": {},
        "state": "failed",
        "created_at": "2020-01-01T00:00:00+00:00",
        "updated_at": "2020-01-01T00:00:00+00:00",
        "retry_count": 2,
        "next_retry_at": "2020-01-01T00:01:00+00:00",
        "last_error": "boom",
        "destination": "central",
    }
    with env.begin(write=True) as txn:
        txn.put(b"rec:legacy", json.dumps(payload).encode())
    env.close()

    backend = LMDBBackend(path=str(tmp_db_path), map_size=10_485_760)
    backend.open()
    stored = backend.get("legacy")
    assert stored is not None
    assert stored.deliveries["central"].state == RecordState.failed
    assert stored.deliveries["central"].retry_count == 2
    assert stored.deliveries["central"].last_error == "boom"
    assert backend.count(RecordState.failed) == 1
    backend.close()