71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
"""In-memory ``NotificationStore`` implementation for tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
from deerflow.persistence.notifications.base import NotificationStore
|
|
|
|
|
|
def _now_iso() -> str:
|
|
return datetime.now(UTC).isoformat()
|
|
|
|
|
|
class MemoryNotificationStore(NotificationStore):
|
|
def __init__(self) -> None:
|
|
self._items: dict[str, dict[str, Any]] = {}
|
|
|
|
async def create(self, data: dict[str, Any]) -> dict[str, Any]:
|
|
record = {
|
|
"id": data["id"],
|
|
"user_id": data["user_id"],
|
|
"type": data["type"],
|
|
"title": data.get("title") or "",
|
|
"body": data.get("body") or "",
|
|
"payload": data.get("payload") or {},
|
|
"read": False,
|
|
"created_at": _now_iso(),
|
|
}
|
|
self._items[record["id"]] = record
|
|
return copy.deepcopy(record)
|
|
|
|
async def list_for_user(
|
|
self,
|
|
user_id: str,
|
|
*,
|
|
unread_only: bool = False,
|
|
limit: int = 50,
|
|
) -> list[dict[str, Any]]:
|
|
filtered = [r for r in self._items.values() if r["user_id"] == user_id]
|
|
if unread_only:
|
|
filtered = [r for r in filtered if not r["read"]]
|
|
filtered.sort(key=lambda r: r["created_at"], reverse=True)
|
|
return [copy.deepcopy(r) for r in filtered[:limit]]
|
|
|
|
async def unread_count(self, user_id: str) -> int:
|
|
return sum(1 for r in self._items.values() if r["user_id"] == user_id and not r["read"])
|
|
|
|
async def mark_read(self, notification_id: str, user_id: str) -> bool:
|
|
record = self._items.get(notification_id)
|
|
if record is None or record["user_id"] != user_id:
|
|
return False
|
|
record["read"] = True
|
|
return True
|
|
|
|
async def mark_all_read(self, user_id: str) -> int:
|
|
updated = 0
|
|
for record in self._items.values():
|
|
if record["user_id"] == user_id and not record["read"]:
|
|
record["read"] = True
|
|
updated += 1
|
|
return updated
|
|
|
|
async def delete(self, notification_id: str, user_id: str) -> bool:
|
|
record = self._items.get(notification_id)
|
|
if record is None or record["user_id"] != user_id:
|
|
return False
|
|
del self._items[notification_id]
|
|
return True
|