feat(notifications): add redis-backed sse notification streaming
This commit is contained in:
200
apps/notifications/tests/test_services.py
Normal file
200
apps/notifications/tests/test_services.py
Normal file
@@ -0,0 +1,200 @@
|
||||
import json
|
||||
from collections import defaultdict
|
||||
|
||||
import pytest
|
||||
|
||||
from apps.notifications import services
|
||||
from apps.notifications.services import RedisNotificationStore
|
||||
|
||||
|
||||
class FakePipeline:
|
||||
def __init__(self, client):
|
||||
self.client = client
|
||||
self.operations = []
|
||||
|
||||
def __getattr__(self, name):
|
||||
def wrapper(*args, **kwargs):
|
||||
self.operations.append((name, args, kwargs))
|
||||
return self
|
||||
|
||||
return wrapper
|
||||
|
||||
def execute(self):
|
||||
results = []
|
||||
for name, args, kwargs in self.operations:
|
||||
results.append(getattr(self.client, name)(*args, **kwargs))
|
||||
self.operations.clear()
|
||||
return results
|
||||
|
||||
|
||||
class FakePubSub:
|
||||
def __init__(self):
|
||||
self.channels = []
|
||||
self.messages = []
|
||||
self.closed = False
|
||||
|
||||
def subscribe(self, channel):
|
||||
self.channels.append(channel)
|
||||
|
||||
def unsubscribe(self, channel):
|
||||
if channel in self.channels:
|
||||
self.channels.remove(channel)
|
||||
|
||||
def get_message(self, timeout=1.0):
|
||||
if self.messages:
|
||||
return self.messages.pop(0)
|
||||
return None
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class FakeRedis:
|
||||
def __init__(self):
|
||||
self.sorted_sets = defaultdict(dict)
|
||||
self.hashes = defaultdict(dict)
|
||||
self.sets = defaultdict(set)
|
||||
self.published = []
|
||||
self.pubsub_instance = FakePubSub()
|
||||
|
||||
def pipeline(self):
|
||||
return FakePipeline(self)
|
||||
|
||||
def zadd(self, key, mapping):
|
||||
self.sorted_sets[key].update(mapping)
|
||||
return len(mapping)
|
||||
|
||||
def hset(self, key, field, value):
|
||||
self.hashes[key][field] = value
|
||||
return 1
|
||||
|
||||
def sadd(self, key, *members):
|
||||
before = len(self.sets[key])
|
||||
self.sets[key].update(members)
|
||||
return len(self.sets[key]) - before
|
||||
|
||||
def zrevrange(self, key, start, stop):
|
||||
items = sorted(
|
||||
self.sorted_sets[key].items(),
|
||||
key=lambda item: (item[1], item[0]),
|
||||
reverse=True,
|
||||
)
|
||||
if stop == -1:
|
||||
return [member for member, _ in items[start:]]
|
||||
return [member for member, _ in items[start : stop + 1]]
|
||||
|
||||
def hget(self, key, field):
|
||||
return self.hashes[key].get(field)
|
||||
|
||||
def zrem(self, key, *members):
|
||||
removed = 0
|
||||
for member in members:
|
||||
if member in self.sorted_sets[key]:
|
||||
del self.sorted_sets[key][member]
|
||||
removed += 1
|
||||
return removed
|
||||
|
||||
def hdel(self, key, *fields):
|
||||
removed = 0
|
||||
for field in fields:
|
||||
if field in self.hashes[key]:
|
||||
del self.hashes[key][field]
|
||||
removed += 1
|
||||
return removed
|
||||
|
||||
def smembers(self, key):
|
||||
return set(self.sets[key])
|
||||
|
||||
def srem(self, key, member):
|
||||
if member in self.sets[key]:
|
||||
self.sets[key].remove(member)
|
||||
return 1
|
||||
return 0
|
||||
|
||||
def zrangebyscore(self, key, min_score, max_score):
|
||||
lower = float("-inf") if min_score == "-inf" else float(min_score)
|
||||
upper = float(max_score)
|
||||
return [
|
||||
member
|
||||
for member, score in self.sorted_sets[key].items()
|
||||
if lower <= score <= upper
|
||||
]
|
||||
|
||||
def zcard(self, key):
|
||||
return len(self.sorted_sets[key])
|
||||
|
||||
def publish(self, channel, message):
|
||||
self.published.append((channel, json.loads(message)))
|
||||
return 1
|
||||
|
||||
def pubsub(self, ignore_subscribe_messages=True):
|
||||
return self.pubsub_instance
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def fake_redis(monkeypatch):
|
||||
redis = FakeRedis()
|
||||
monkeypatch.setattr(services, "redis_client", redis)
|
||||
return redis
|
||||
|
||||
|
||||
def test_add_publishes_notification_and_unread_count(fake_redis, settings):
|
||||
settings.NOTIFICATIONS_ENABLED = True
|
||||
|
||||
notification = RedisNotificationStore.add(
|
||||
"user-1",
|
||||
{
|
||||
"title": "Build finished",
|
||||
"message": "Your deploy completed.",
|
||||
"level": "success",
|
||||
},
|
||||
)
|
||||
|
||||
assert notification["title"] == "Build finished"
|
||||
assert notification["message"] == "Your deploy completed."
|
||||
assert notification["level"] == "success"
|
||||
assert len(fake_redis.published) == 2
|
||||
channel, payload = fake_redis.published[0]
|
||||
assert channel == f"{settings.NOTIFICATION_REDIS_CHANNEL_PREFIX}:user-1"
|
||||
assert payload["event"] == "notification"
|
||||
assert payload["data"]["notification"]["id"] == notification["id"]
|
||||
assert payload["data"]["unread_count"] == 1
|
||||
|
||||
|
||||
def test_mark_seen_and_mark_all_seen_publish_sync_events(fake_redis, settings):
|
||||
settings.NOTIFICATIONS_ENABLED = True
|
||||
first = RedisNotificationStore.add("user-2", {"title": "First"})
|
||||
second = RedisNotificationStore.add("user-2", {"title": "Second"})
|
||||
fake_redis.published.clear()
|
||||
|
||||
payload = RedisNotificationStore.mark_seen("user-2", first["id"])
|
||||
|
||||
assert payload["notification_id"] == first["id"]
|
||||
assert payload["deleted"] is False
|
||||
assert payload["notification"]["is_seen"] is True
|
||||
assert fake_redis.published[0][1]["event"] == "notification_seen"
|
||||
|
||||
fake_redis.published.clear()
|
||||
updated = RedisNotificationStore.mark_all_seen("user-2")
|
||||
|
||||
assert updated == 2
|
||||
assert fake_redis.published[0][1]["event"] == "notification_mark_all_read"
|
||||
assert fake_redis.published[1][1]["event"] == "unread_count"
|
||||
assert fake_redis.published[1][1]["data"]["unread_count"] == 0
|
||||
|
||||
|
||||
def test_list_returns_total_count_and_filtered_notifications(fake_redis):
|
||||
RedisNotificationStore.add("user-3", {"title": "General", "type": "general"})
|
||||
RedisNotificationStore.add("user-3", {"title": "Billing", "type": "billing"})
|
||||
RedisNotificationStore.add("user-3", {"title": "General 2", "type": "general"})
|
||||
|
||||
notifications, total_count = RedisNotificationStore.list(
|
||||
"user-3",
|
||||
limit=1,
|
||||
offset=0,
|
||||
type_filter="general",
|
||||
)
|
||||
|
||||
assert total_count == 2
|
||||
assert len(notifications) == 1
|
||||
assert notifications[0]["type"] == "general"
|
||||
165
apps/notifications/tests/test_views.py
Normal file
165
apps/notifications/tests/test_views.py
Normal file
@@ -0,0 +1,165 @@
|
||||
import json
|
||||
import time
|
||||
from datetime import timedelta
|
||||
|
||||
import pytest
|
||||
from django.utils import timezone
|
||||
from rest_framework.test import APIClient
|
||||
|
||||
from apps.notifications import services, views
|
||||
from apps.notifications.services import RedisNotificationStore
|
||||
from apps.notifications.tests.test_services import FakePubSub, FakeRedis
|
||||
from apps.users.models import User
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def fake_redis(monkeypatch):
|
||||
redis = FakeRedis()
|
||||
monkeypatch.setattr(services, "redis_client", redis)
|
||||
return redis
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def user(db):
|
||||
return User.objects.create_user(mobile="09121111111", password="secret123")
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def second_user(db):
|
||||
return User.objects.create_user(mobile="09122222222", password="secret123")
|
||||
|
||||
|
||||
def _read_sse_chunks(response, count):
|
||||
iterator = iter(response.streaming_content)
|
||||
chunks = []
|
||||
for _ in range(count):
|
||||
chunk = next(iterator)
|
||||
if isinstance(chunk, bytes):
|
||||
chunk = chunk.decode("utf-8")
|
||||
chunks.append(chunk)
|
||||
response.close()
|
||||
return chunks
|
||||
|
||||
|
||||
def _parse_sse_data(chunk: str) -> dict:
|
||||
for line in chunk.splitlines():
|
||||
if line.startswith("data: "):
|
||||
return json.loads(line.removeprefix("data: "))
|
||||
raise AssertionError("SSE payload did not include data")
|
||||
|
||||
|
||||
def test_stream_token_endpoint_returns_short_lived_token(user):
|
||||
client = APIClient()
|
||||
client.force_authenticate(user=user)
|
||||
|
||||
response = client.post("/api/notifications/stream-token/")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.data["token"]
|
||||
assert response.data["expires_in"] > 0
|
||||
|
||||
|
||||
def test_stream_endpoint_rejects_missing_and_expired_token(user, settings):
|
||||
client = APIClient()
|
||||
|
||||
missing = client.get("/api/notifications/stream/")
|
||||
assert missing.status_code == 401
|
||||
|
||||
settings.NOTIFICATION_STREAM_TOKEN_LIFETIME_SECONDS = 1
|
||||
token = views._issue_stream_token_for_user(str(user.id))
|
||||
time.sleep(1.1)
|
||||
|
||||
expired = client.get(f"/api/notifications/stream/?token={token}")
|
||||
assert expired.status_code == 401
|
||||
|
||||
|
||||
def test_stream_endpoint_sends_only_current_users_notifications(
|
||||
fake_redis, user, second_user, monkeypatch
|
||||
):
|
||||
RedisNotificationStore.add(str(user.id), {"title": "For current user"})
|
||||
RedisNotificationStore.add(str(second_user.id), {"title": "For another user"})
|
||||
pubsub = FakePubSub()
|
||||
monkeypatch.setattr(RedisNotificationStore, "get_pubsub", classmethod(lambda cls: pubsub))
|
||||
token = views._issue_stream_token_for_user(str(user.id))
|
||||
|
||||
client = APIClient()
|
||||
response = client.get(
|
||||
f"/api/notifications/stream/?token={token}",
|
||||
HTTP_ACCEPT="text/event-stream",
|
||||
)
|
||||
retry_line, connected_chunk = _read_sse_chunks(response, 2)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert retry_line.startswith("retry:")
|
||||
connected = _parse_sse_data(connected_chunk)
|
||||
assert connected["unread_count"] == 1
|
||||
assert [item["title"] for item in connected["notifications"]] == ["For current user"]
|
||||
|
||||
|
||||
def test_stream_endpoint_emits_heartbeat(fake_redis, user, settings, monkeypatch):
|
||||
pubsub = FakePubSub()
|
||||
monkeypatch.setattr(RedisNotificationStore, "get_pubsub", classmethod(lambda cls: pubsub))
|
||||
settings.NOTIFICATION_SSE_HEARTBEAT_SECONDS = 1
|
||||
|
||||
first_now = timezone.now()
|
||||
tick_values = iter(
|
||||
[
|
||||
first_now,
|
||||
first_now,
|
||||
first_now + timedelta(seconds=2),
|
||||
first_now + timedelta(seconds=2),
|
||||
first_now + timedelta(seconds=2),
|
||||
first_now + timedelta(seconds=2),
|
||||
]
|
||||
)
|
||||
last_tick = first_now + timedelta(seconds=2)
|
||||
|
||||
def fake_now():
|
||||
return next(tick_values, last_tick)
|
||||
|
||||
monkeypatch.setattr(views.timezone, "now", fake_now)
|
||||
view = views.NotificationStreamView()
|
||||
stream = view._build_stream(str(user.id))
|
||||
|
||||
chunks = [next(stream) for _ in range(4)]
|
||||
stream.close()
|
||||
|
||||
assert "event: ping" in chunks[3]
|
||||
|
||||
|
||||
def test_notification_list_and_seen_endpoints_work(fake_redis, user):
|
||||
notification = RedisNotificationStore.add(
|
||||
str(user.id),
|
||||
{"title": "Deploy succeeded", "type": "deploy"},
|
||||
)
|
||||
|
||||
client = APIClient()
|
||||
client.force_authenticate(user=user)
|
||||
|
||||
list_response = client.get("/api/notifications/list/?type=deploy")
|
||||
assert list_response.status_code == 200
|
||||
assert list_response.data["count"] == 1
|
||||
assert list_response.data["unread_count"] == 1
|
||||
assert list_response.data["notifications"][0]["title"] == "Deploy succeeded"
|
||||
|
||||
seen_response = client.post("/api/notifications/seen/", {"id": notification["id"]}, format="json")
|
||||
assert seen_response.status_code == 200
|
||||
assert seen_response.data["marked_read"] is True
|
||||
assert seen_response.data["notification"]["is_seen"] is True
|
||||
|
||||
|
||||
def test_notification_delete_endpoint_removes_notification(fake_redis, user):
|
||||
notification = RedisNotificationStore.add(
|
||||
str(user.id),
|
||||
{"title": "Delete me", "type": "deploy"},
|
||||
)
|
||||
|
||||
client = APIClient()
|
||||
client.force_authenticate(user=user)
|
||||
|
||||
response = client.delete(f"/api/notifications/{notification['id']}/")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.data["deleted"] is True
|
||||
assert response.data["notification_id"] == notification["id"]
|
||||
assert RedisNotificationStore.get(str(user.id), notification["id"]) is None
|
||||
Reference in New Issue
Block a user