server-26/drb-c2-core/tests/fake_firestore.py

"""
In-memory Firestore stand-in for the Discord lease/stream tests.

Models the one property those modules depend on: a transaction is
optimistic — every document (and every collection listing) it read is
re-checked at commit, and if anything changed underneath it the whole
function is re-run, exactly like the real SDK's ``transactional`` retry. A
``before_commit`` hook lets a test hold two transactions at the commit point
together, so a race the code must survive actually happens in the test.

Install with ``FakeFirestore.install(monkeypatch)``; it patches
``app.internal.firestore``'s db + async helpers in place.
"""
from __future__ import annotations

import asyncio
import copy
import threading
from typing import Callable, Optional


class Conflict(Exception):
    pass


class _Snap:
    def __init__(self, doc_id: str, data: Optional[dict]):
        self.id = doc_id
        self._data = data
        self.exists = data is not None

    def to_dict(self):
        return copy.deepcopy(self._data) if self._data is not None else None


class _DocRef:
    def __init__(self, db: "FakeFirestore", coll: str, doc_id: str):
        self._db, self.coll, self.id = db, coll, doc_id

    def get(self, transaction=None):
        with self._db.lock:
            data = self._db.data.get(self.coll, {}).get(self.id)
            ver = self._db.versions.get((self.coll, self.id), 0)
        if transaction is not None:
            transaction.reads[("doc", self.coll, self.id)] = ver
        return _Snap(self.id, copy.deepcopy(data))


class _Collection:
    def __init__(self, db: "FakeFirestore", name: str):
        self._db, self.name = db, name

    def document(self, doc_id: str) -> _DocRef:
        return _DocRef(self._db, self.name, doc_id)

    def stream(self, transaction=None):
        with self._db.lock:
            docs = {k: copy.deepcopy(v) for k, v in self._db.data.get(self.name, {}).items()}
            vers = {k: self._db.versions.get((self.name, k), 0) for k in docs}
        if transaction is not None:
            transaction.reads[("coll", self.name)] = frozenset(docs)
            for k, v in vers.items():
                transaction.reads[("doc", self.name, k)] = v
        return [_Snap(k, v) for k, v in docs.items()]


class _Txn:
    def __init__(self):
        self.reads: dict = {}
        self.writes: list = []

    def set(self, ref: _DocRef, data: dict, merge: bool = False):
        self.writes.append(("set", ref.coll, ref.id, copy.deepcopy(data), merge))

    def update(self, ref: _DocRef, data: dict):
        self.writes.append(("update", ref.coll, ref.id, copy.deepcopy(data), True))

    def delete(self, ref: _DocRef):
        self.writes.append(("delete", ref.coll, ref.id, {}, False))


class FakeFirestore:
    def __init__(self):
        self.data: dict[str, dict[str, dict]] = {}
        self.versions: dict[tuple, int] = {}
        self.lock = threading.RLock()
        self.before_commit: Optional[Callable[[int], None]] = None
        self.commits = 0
        self.retries = 0

    # --- sync surface used inside transactions ---------------------------
    def collection(self, name: str) -> _Collection:
        return _Collection(self, name)

    def _write(self, kind: str, coll: str, doc_id: str, data: dict, merge: bool):
        docs = self.data.setdefault(coll, {})
        if kind == "delete":
            docs.pop(doc_id, None)
            self.versions[(coll, doc_id)] = self.versions.get((coll, doc_id), 0) + 1
            return
        if kind == "update" and doc_id not in docs:
            raise KeyError(f"No document to update: {coll}/{doc_id}")
        if merge and doc_id in docs:
            docs[doc_id] = {**docs[doc_id], **data}
        else:
            docs[doc_id] = dict(data)
        self.versions[(coll, doc_id)] = self.versions.get((coll, doc_id), 0) + 1

    def run_transaction_sync(self, fn, max_attempts: int = 20):
        for attempt in range(max_attempts):
            txn = _Txn()
            result = fn(txn)
            if self.before_commit is not None:
                self.before_commit(attempt)
            with self.lock:
                stale = False
                for key, seen in txn.reads.items():
                    if key[0] == "doc":
                        if self.versions.get((key[1], key[2]), 0) != seen:
                            stale = True
                            break
                    elif frozenset(self.data.get(key[1], {})) != seen:
                        stale = True
                        break
                if not stale:
                    for w in txn.writes:
                        self._write(*w)
                    self.commits += 1
                    return result
            self.retries += 1
        raise Conflict("transaction kept conflicting")

    # --- async helpers mirroring app.internal.firestore ------------------
    async def run_transaction(self, fn):
        return await asyncio.to_thread(self.run_transaction_sync, fn)

    async def doc_get(self, coll: str, doc_id: str):
        with self.lock:
            d = self.data.get(coll, {}).get(doc_id)
            return copy.deepcopy(d) if d is not None else None

    async def doc_get_cached(self, coll: str, doc_id: str, ttl: float = 300.0):
        return await self.doc_get(coll, doc_id)

    async def doc_set(self, coll: str, doc_id: str, data: dict, merge: bool = True):
        with self.lock:
            self._write("set", coll, doc_id, copy.deepcopy(data), merge)

    async def doc_update(self, coll: str, doc_id: str, data: dict):
        with self.lock:
            self._write("update", coll, doc_id, copy.deepcopy(data), True)

    async def doc_delete(self, coll: str, doc_id: str):
        with self.lock:
            self.data.get(coll, {}).pop(doc_id, None)
            self.versions[(coll, doc_id)] = self.versions.get((coll, doc_id), 0) + 1

    async def collection_list(self, coll: str, **filters):
        with self.lock:
            docs = [copy.deepcopy(d) for d in self.data.get(coll, {}).values()]
        return [d for d in docs if all(d.get(k) == v for k, v in filters.items())]

    # --- test conveniences ------------------------------------------------
    def put(self, coll: str, doc_id: str, data: dict):
        with self.lock:
            self._write("set", coll, doc_id, copy.deepcopy(data), False)

    def get(self, coll: str, doc_id: str) -> Optional[dict]:
        return copy.deepcopy(self.data.get(coll, {}).get(doc_id))

    @classmethod
    def install(cls, monkeypatch) -> "FakeFirestore":
        from app.internal import firestore as fstore
        fake = cls()
        monkeypatch.setattr(fstore, "db", fake)
        for name in ("run_transaction", "doc_get", "doc_get_cached", "doc_set", "doc_update",
                     "doc_delete", "collection_list"):
            monkeypatch.setattr(fstore, name, getattr(fake, name))
        return fake