""" 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