ansible-collection-famedly-matrix/tests/unit/mock_nio/utils/RoomSimulator.py

import json
import secrets
import string
from copy import deepcopy
from typing import Any, Optional, Union
from typing import Dict


class RoomEvents:
    M_ROOM_CREATE = {
        "content": {
            "creator": "@example:localhost",
        },
        "sender": "@example:localhost",
        "state_key": "",
        "type": "m.room.create",
    }
    M_ROOM_MEMBER = {
        "content": {"membership": "join"},
        "membership": "join",
        "sender": "@example:localhost",
        "state_key": "@example:localhost",
        "type": "m.room.member",
    }
    # This is not to spec!
    M_ROOM_DUMMY = {
        "content": {"dummy": "I'm a test message"},
        "sender": "@example:localhost",
        "state_key": "@example:localhost",
        "type": "m.room.dummy",
    }


class RoomSimulator:
    INITIAL_DATA = {"rooms": {}, "room_directory": {}}

    def __init__(self, data: str = json.dumps(INITIAL_DATA)):
        try:
            self.data = json.loads(data)
        except (json.JSONDecodeError, TypeError):
            self.data = self.INITIAL_DATA

    def random_event_id(self):
        return "".join(
            secrets.choice(string.ascii_uppercase + string.digits) for i in range(64)
        )

    def send_event(self, room_id: str, event: Dict[str, Any]) -> Dict[str, Any]:
        event["event_id"] = self.random_event_id()
        for e in self.get_events(room_id):
            if e["type"] == event["type"] and e["state_key"] == event["state_key"]:
                event["replaces_state"] = e["event_id"]
        self.add_event(room_id=room_id, event=event)
        return event

    def get_events(self, room_id: str) -> Dict[str, Any]:
        return self.data["rooms"][room_id]

    def add_event(self, room_id: str, event: Dict[str, Any]) -> Dict[str, Any]:
        return self.data["rooms"][room_id].append(event)

    def get_state_event(
        self, room_id: str, event_type: str, state_key: str
    ) -> Union[None, Dict[str, Any]]:
        events = list[Dict[str, Any]]()
        for e in self.get_events(room_id):
            if e["type"] == event_type and e["state_key"] == state_key:
                events.append(e)
        if len(events) > 0:
            return events[-1]
        return None

    def add_room(self, room_id: str, alias: Optional[str] = None):
        self.data["rooms"][room_id] = []
        if alias is not None:
            self.data["room_directory"][alias] = room_id
        return room_id

    def resolve_alias(self, alias: str) -> str:
        return self.data["room_directory"].get(alias, None)

    def export(self) -> str:
        return json.dumps(self.data)

    # Matrix Tasks

    def create_room(self, alias: Optional[str] = None):
        room_id = (
            "!"
            + "".join(
                secrets.choice(string.ascii_uppercase + string.digits)
                for i in range(10)
            )
            + ":matrix.example.tld"
        )
        self.add_room(room_id, alias)
        return room_id

    # Matrix Events

    def m_room_create(self, room_id: str, creator: str) -> Dict[str, Any]:
        data = deepcopy(RoomEvents.M_ROOM_CREATE)
        data["content"]["creator"] = creator
        data["sender"] = creator
        return self.send_event(room_id=room_id, event=data)

    def m_room_member(
        self, room_id: str, membership: str, sender: str, state_key: str
    ) -> Dict[str, Any]:
        data = deepcopy(RoomEvents.M_ROOM_MEMBER)
        data["content"]["membership"] = membership
        data["membership"] = membership
        data["sender"] = sender
        data["state_key"] = state_key
        return self.send_event(room_id=room_id, event=data)