How to Mock Redis Streams Consumer Groups in Python

import time
import json
from collections import defaultdict

class RedisStreamMock:
    def __init__(self):
        self.streams = defaultdict(list)
        self.consumer_groups = defaultdict(dict)
        self.pending_entries = defaultdict(list)

    def xadd(self, stream, fields):
        entry_id = f"{time.time_ns()}-{len(self.streams[stream])}"
        entry = {**fields, "_id": entry_id}
        self.streams[stream].append(entry)
        return entry_id

    def xgroup_create(self, stream, group, start="0"):
        self.consumer_groups[stream][group] = {"last_id": start, "consumers": set()}

    def xreadgroup(self, group, consumer, stream, count=1):
        group_data = self.consumer_groups[stream][group]
        group_data["consumers"].add(consumer)
        last_id = group_data["last_id"]
        available = [e for e in self.streams[stream] if e["_id"] > last_id]
        if not available:
            return []
        selected = available[:count]
        group_data["last_id"] = selected[-1]["_id"]
        for entry in selected:
            self.pending_entries[stream].append({"group": group, "consumer": consumer, "entry": entry})
        return [{"stream": stream, "entries": selected}]

    def xack(self, stream, group, entry_id):
        self.pending_entries[stream] = [
            p for p in self.pending_entries[stream]
            if not (p["group"] == group and p["entry"]["_id"] == entry_id)
        ]
        return True


if __name__ == "__main__":
    redis = RedisStreamMock()

    # Producer adds two messages
    redis.xadd("orders", {"user": "alice", "amount": 100})
    redis.xadd("orders", {"user": "bob", "amount": 250})

    # Create a consumer group
    redis.xgroup_create("orders", "payments")

    # Consumer reads the first message
    result = redis.xreadgroup("payments", "worker-1", "orders", count=1)
    first_entry = result[0]["entries"][0]
    print("Read:", first_entry["user"], first_entry["amount"], first_entry["_id"])

    # Acknowledge it
    redis.xack("orders", "payments", first_entry["_id"])

    # Second consumer reads the next message
    result2 = redis.xreadgroup("payments", "worker-2", "orders", count=1)
    second_entry = result2[0]["entries"][0]
    print("Read:", second_entry["user"], second_entry["amount"], second_entry["_id"])

    # No more messages
    empty = redis.xreadgroup("payments", "worker-2", "orders", count=1)
    print("Empty read:", empty == [])