def test_prefilled_cache(self):
     """
     Providing a prefilled cache to StreamChangeCache will result in a cache
     with the prefilled-cache entered in.
     """
     cache = StreamChangeCache("#test", 1, prefilled_cache={"[email protected]": 2})
     self.assertTrue(cache.has_entity_changed("[email protected]", 1))

class SlavedGroupServerStore(BaseSlavedStore):
    def __init__(self, db_conn, hs):
        super(SlavedGroupServerStore, self).__init__(db_conn, hs)

        self.hs = hs

        self._group_updates_id_gen = SlavedIdTracker(
            db_conn, "local_group_updates", "stream_id",
        )
        self._group_updates_stream_cache = StreamChangeCache(
            "_group_updates_stream_cache", self._group_updates_id_gen.get_current_token(),
        )

    get_groups_changes_for_user = DataStore.get_groups_changes_for_user.__func__
    get_group_stream_token = DataStore.get_group_stream_token.__func__
    get_all_groups_for_user = DataStore.get_all_groups_for_user.__func__

    def stream_positions(self):
        result = super(SlavedGroupServerStore, self).stream_positions()
        result["groups"] = self._group_updates_id_gen.get_current_token()
        return result

    def process_replication_rows(self, stream_name, token, rows):
        if stream_name == "groups":
            self._group_updates_id_gen.advance(token)
            for row in rows:
                self._group_updates_stream_cache.entity_has_changed(
                    row.user_id, token
                )

        return super(SlavedGroupServerStore, self).process_replication_rows(
            stream_name, token, rows
        )

    def __init__(self, db_conn, hs):
        super(SlavedEventStore, self).__init__(db_conn, hs)
        self._stream_id_gen = SlavedIdTracker(
            db_conn, "events", "stream_ordering",
        )
        self._backfill_id_gen = SlavedIdTracker(
            db_conn, "events", "stream_ordering", step=-1
        )
        events_max = self._stream_id_gen.get_current_token()
        event_cache_prefill, min_event_val = self._get_cache_dict(
            db_conn, "events",
            entity_column="room_id",
            stream_column="stream_ordering",
            max_value=events_max,
        )
        self._events_stream_cache = StreamChangeCache(
            "EventsRoomStreamChangeCache", min_event_val,
            prefilled_cache=event_cache_prefill,
        )
        self._membership_stream_cache = StreamChangeCache(
            "MembershipStreamChangeCache", events_max,
        )

        self.stream_ordering_month_ago = 0
        self._stream_order_on_start = self.get_room_max_stream_ordering()

class SlavedReceiptsStore(BaseSlavedStore):

    def __init__(self, db_conn, hs):
        super(SlavedReceiptsStore, self).__init__(db_conn, hs)

        self._receipts_id_gen = SlavedIdTracker(
            db_conn, "receipts_linearized", "stream_id"
        )

        self._receipts_stream_cache = StreamChangeCache(
            "ReceiptsRoomChangeCache", self._receipts_id_gen.get_current_token()
        )

    get_receipts_for_user = ReceiptsStore.__dict__["get_receipts_for_user"]
    get_linearized_receipts_for_room = (
        ReceiptsStore.__dict__["get_linearized_receipts_for_room"]
    )
    _get_linearized_receipts_for_rooms = (
        ReceiptsStore.__dict__["_get_linearized_receipts_for_rooms"]
    )
    get_last_receipt_event_id_for_user = (
        ReceiptsStore.__dict__["get_last_receipt_event_id_for_user"]
    )

    get_max_receipt_stream_id = DataStore.get_max_receipt_stream_id.__func__
    get_all_updated_receipts = DataStore.get_all_updated_receipts.__func__

    get_linearized_receipts_for_rooms = (
        DataStore.get_linearized_receipts_for_rooms.__func__
    )

    def stream_positions(self):
        result = super(SlavedReceiptsStore, self).stream_positions()
        result["receipts"] = self._receipts_id_gen.get_current_token()
        return result

    def process_replication(self, result):
        stream = result.get("receipts")
        if stream:
            self._receipts_id_gen.advance(int(stream["position"]))
            for row in stream["rows"]:
                position, room_id, receipt_type, user_id = row[:4]
                self.invalidate_caches_for_receipt(room_id, receipt_type, user_id)
                self._receipts_stream_cache.entity_has_changed(room_id, position)

        return super(SlavedReceiptsStore, self).process_replication(result)

    def invalidate_caches_for_receipt(self, room_id, receipt_type, user_id):
        self.get_receipts_for_user.invalidate((user_id, receipt_type))
        self.get_linearized_receipts_for_room.invalidate_many((room_id,))
        self.get_last_receipt_event_id_for_user.invalidate(
            (user_id, room_id, receipt_type)
        )

    def __init__(self, db_conn, hs):
        super(SlavedDeviceStore, self).__init__(db_conn, hs)

        self.hs = hs

        self._device_list_id_gen = SlavedIdTracker(
            db_conn, "device_lists_stream", "stream_id",
        )
        device_list_max = self._device_list_id_gen.get_current_token()
        self._device_list_stream_cache = StreamChangeCache(
            "DeviceListStreamChangeCache", device_list_max,
        )
        self._device_list_federation_stream_cache = StreamChangeCache(
            "DeviceListFederationStreamChangeCache", device_list_max,
        )

class SlavedDeviceInboxStore(BaseSlavedStore):
    def __init__(self, db_conn, hs):
        super(SlavedDeviceInboxStore, self).__init__(db_conn, hs)
        self._device_inbox_id_gen = SlavedIdTracker(
            db_conn, "device_max_stream_id", "stream_id",
        )
        self._device_inbox_stream_cache = StreamChangeCache(
            "DeviceInboxStreamChangeCache",
            self._device_inbox_id_gen.get_current_token()
        )
        self._device_federation_outbox_stream_cache = StreamChangeCache(
            "DeviceFederationOutboxStreamChangeCache",
            self._device_inbox_id_gen.get_current_token()
        )

        self._last_device_delete_cache = ExpiringCache(
            cache_name="last_device_delete_cache",
            clock=self._clock,
            max_len=10000,
            expiry_ms=30 * 60 * 1000,
        )

    get_to_device_stream_token = DataStore.get_to_device_stream_token.__func__
    get_new_messages_for_device = DataStore.get_new_messages_for_device.__func__
    get_new_device_msgs_for_remote = DataStore.get_new_device_msgs_for_remote.__func__
    delete_messages_for_device = DataStore.delete_messages_for_device.__func__
    delete_device_msgs_for_remote = DataStore.delete_device_msgs_for_remote.__func__

    def stream_positions(self):
        result = super(SlavedDeviceInboxStore, self).stream_positions()
        result["to_device"] = self._device_inbox_id_gen.get_current_token()
        return result

    def process_replication_rows(self, stream_name, token, rows):
        if stream_name == "to_device":
            self._device_inbox_id_gen.advance(token)
            for row in rows:
                if row.entity.startswith("@"):
                    self._device_inbox_stream_cache.entity_has_changed(
                        row.entity, token
                    )
                else:
                    self._device_federation_outbox_stream_cache.entity_has_changed(
                        row.entity, token
                    )
        return super(SlavedDeviceInboxStore, self).process_replication_rows(
            stream_name, token, rows
        )

class SlavedDeviceStore(BaseSlavedStore):
    def __init__(self, db_conn, hs):
        super(SlavedDeviceStore, self).__init__(db_conn, hs)

        self.hs = hs

        self._device_list_id_gen = SlavedIdTracker(
            db_conn, "device_lists_stream", "stream_id",
        )
        device_list_max = self._device_list_id_gen.get_current_token()
        self._device_list_stream_cache = StreamChangeCache(
            "DeviceListStreamChangeCache", device_list_max,
        )
        self._device_list_federation_stream_cache = StreamChangeCache(
            "DeviceListFederationStreamChangeCache", device_list_max,
        )

    get_device_stream_token = __func__(DataStore.get_device_stream_token)
    get_user_whose_devices_changed = __func__(DataStore.get_user_whose_devices_changed)
    get_devices_by_remote = __func__(DataStore.get_devices_by_remote)
    _get_devices_by_remote_txn = __func__(DataStore._get_devices_by_remote_txn)
    _get_e2e_device_keys_txn = __func__(DataStore._get_e2e_device_keys_txn)
    mark_as_sent_devices_by_remote = __func__(DataStore.mark_as_sent_devices_by_remote)
    _mark_as_sent_devices_by_remote_txn = (
        __func__(DataStore._mark_as_sent_devices_by_remote_txn)
    )
    count_e2e_one_time_keys = EndToEndKeyStore.__dict__["count_e2e_one_time_keys"]

    def stream_positions(self):
        result = super(SlavedDeviceStore, self).stream_positions()
        result["device_lists"] = self._device_list_id_gen.get_current_token()
        return result

    def process_replication_rows(self, stream_name, token, rows):
        if stream_name == "device_lists":
            self._device_list_id_gen.advance(token)
            for row in rows:
                self._device_list_stream_cache.entity_has_changed(
                    row.user_id, token
                )

                if row.destination:
                    self._device_list_federation_stream_cache.entity_has_changed(
                        row.destination, token
                    )
        return super(SlavedDeviceStore, self).process_replication_rows(
            stream_name, token, rows
        )

    def __init__(self, db_conn, hs):
        account_max = self.get_max_account_data_stream_id()
        self._account_data_stream_cache = StreamChangeCache(
            "AccountDataAndTagsChangeCache", account_max,
        )

        super(AccountDataWorkerStore, self).__init__(db_conn, hs)

    def __init__(self, hs):
        self.store = hs.get_datastore()
        self.server_name = hs.config.server_name
        self.auth = hs.get_auth()
        self.is_mine_id = hs.is_mine_id
        self.notifier = hs.get_notifier()
        self.state = hs.get_state_handler()

        self.hs = hs

        self.clock = hs.get_clock()
        self.wheel_timer = WheelTimer(bucket_size=5000)

        self.federation = hs.get_federation_sender()

        hs.get_federation_registry().register_edu_handler("m.typing", self._recv_edu)

        hs.get_distributor().observe("user_left_room", self.user_left_room)

        self._member_typing_until = {}  # clock time we expect to stop
        self._member_last_federation_poke = {}

        self._latest_room_serial = 0
        self._reset()

        # caches which room_ids changed at which serials
        self._typing_stream_change_cache = StreamChangeCache(
            "TypingStreamChangeCache", self._latest_room_serial,
        )

        self.clock.looping_call(
            self._handle_timeouts,
            5000,
        )

class SlavedDeviceStore(EndToEndKeyWorkerStore, DeviceWorkerStore, BaseSlavedStore):
    def __init__(self, db_conn, hs):
        super(SlavedDeviceStore, self).__init__(db_conn, hs)

        self.hs = hs

        self._device_list_id_gen = SlavedIdTracker(
            db_conn, "device_lists_stream", "stream_id",
        )
        device_list_max = self._device_list_id_gen.get_current_token()
        self._device_list_stream_cache = StreamChangeCache(
            "DeviceListStreamChangeCache", device_list_max,
        )
        self._device_list_federation_stream_cache = StreamChangeCache(
            "DeviceListFederationStreamChangeCache", device_list_max,
        )

    def stream_positions(self):
        result = super(SlavedDeviceStore, self).stream_positions()
        result["device_lists"] = self._device_list_id_gen.get_current_token()
        return result

    def process_replication_rows(self, stream_name, token, rows):
        if stream_name == "device_lists":
            self._device_list_id_gen.advance(token)
            for row in rows:
                self._invalidate_caches_for_devices(
                    token, row.user_id, row.destination,
                )
        return super(SlavedDeviceStore, self).process_replication_rows(
            stream_name, token, rows
        )

    def _invalidate_caches_for_devices(self, token, user_id, destination):
        self._device_list_stream_cache.entity_has_changed(
            user_id, token
        )

        if destination:
            self._device_list_federation_stream_cache.entity_has_changed(
                destination, token
            )

        self._get_cached_devices_for_user.invalidate((user_id,))
        self._get_cached_user_device.invalidate_many((user_id,))
        self.get_device_list_last_stream_id_for_remote.invalidate((user_id,))

    def test_has_any_entity_changed(self):
        """
        StreamChangeCache.has_any_entity_changed will return True if any
        entities have been changed since the provided stream position, and
        False if they have not.  If the cache has entries and the provided
        stream position is before it, it will return True, otherwise False if
        the cache has no entries.
        """
        cache = StreamChangeCache("#test", 1)

        # With no entities, it returns False for the past, present, and future.
        self.assertFalse(cache.has_any_entity_changed(0))
        self.assertFalse(cache.has_any_entity_changed(1))
        self.assertFalse(cache.has_any_entity_changed(2))

        # We add an entity
        cache.entity_has_changed("[email protected]", 2)

        # With an entity, it returns True for the past, the stream start
        # position, and False for the stream position the entity was changed
        # on and ones after it.
        self.assertTrue(cache.has_any_entity_changed(0))
        self.assertTrue(cache.has_any_entity_changed(1))
        self.assertFalse(cache.has_any_entity_changed(2))
        self.assertFalse(cache.has_any_entity_changed(3))

class SlavedPushRuleStore(SlavedEventStore):
    def __init__(self, db_conn, hs):
        super(SlavedPushRuleStore, self).__init__(db_conn, hs)
        self._push_rules_stream_id_gen = SlavedIdTracker(
            db_conn, "push_rules_stream", "stream_id",
        )
        self.push_rules_stream_cache = StreamChangeCache(
            "PushRulesStreamChangeCache",
            self._push_rules_stream_id_gen.get_current_token(),
        )

    get_push_rules_for_user = PushRuleStore.__dict__["get_push_rules_for_user"]
    get_push_rules_enabled_for_user = (
        PushRuleStore.__dict__["get_push_rules_enabled_for_user"]
    )
    have_push_rules_changed_for_user = (
        DataStore.have_push_rules_changed_for_user.__func__
    )

    def get_push_rules_stream_token(self):
        return (
            self._push_rules_stream_id_gen.get_current_token(),
            self._stream_id_gen.get_current_token(),
        )

    def stream_positions(self):
        result = super(SlavedPushRuleStore, self).stream_positions()
        result["push_rules"] = self._push_rules_stream_id_gen.get_current_token()
        return result

    def process_replication(self, result):
        stream = result.get("push_rules")
        if stream:
            for row in stream["rows"]:
                position = row[0]
                user_id = row[2]
                self.get_push_rules_for_user.invalidate((user_id,))
                self.get_push_rules_enabled_for_user.invalidate((user_id,))
                self.push_rules_stream_cache.entity_has_changed(
                    user_id, position
                )

            self._push_rules_stream_id_gen.advance(int(stream["position"]))

        return super(SlavedPushRuleStore, self).process_replication(result)

 def __init__(self, db_conn, hs):
     super(SlavedAccountDataStore, self).__init__(db_conn, hs)
     self._account_data_id_gen = SlavedIdTracker(
         db_conn, "account_data_max_stream_id", "stream_id",
     )
     self._account_data_stream_cache = StreamChangeCache(
         "AccountDataAndTagsChangeCache",
         self._account_data_id_gen.get_current_token(),
     )

 def __init__(self, db_conn, hs):
     super(SlavedPushRuleStore, self).__init__(db_conn, hs)
     self._push_rules_stream_id_gen = SlavedIdTracker(
         db_conn, "push_rules_stream", "stream_id",
     )
     self.push_rules_stream_cache = StreamChangeCache(
         "PushRulesStreamChangeCache",
         self._push_rules_stream_id_gen.get_current_token(),
     )

    def __init__(self, db_conn, hs):
        super(SlavedReceiptsStore, self).__init__(db_conn, hs)

        self._receipts_id_gen = SlavedIdTracker(
            db_conn, "receipts_linearized", "stream_id"
        )

        self._receipts_stream_cache = StreamChangeCache(
            "ReceiptsRoomChangeCache", self._receipts_id_gen.get_current_token()
        )

    def __init__(self, db_conn, hs):
        super(StreamWorkerStore, self).__init__(db_conn, hs)

        events_max = self.get_room_max_stream_ordering()
        event_cache_prefill, min_event_val = self._get_cache_dict(
            db_conn, "events",
            entity_column="room_id",
            stream_column="stream_ordering",
            max_value=events_max,
        )
        self._events_stream_cache = StreamChangeCache(
            "EventsRoomStreamChangeCache", min_event_val,
            prefilled_cache=event_cache_prefill,
        )
        self._membership_stream_cache = StreamChangeCache(
            "MembershipStreamChangeCache", events_max,
        )

        self._stream_order_on_start = self.get_room_max_stream_ordering()

    def __init__(self, db_conn, hs):
        super(SlavedGroupServerStore, self).__init__(db_conn, hs)

        self.hs = hs

        self._group_updates_id_gen = SlavedIdTracker(
            db_conn, "local_group_updates", "stream_id",
        )
        self._group_updates_stream_cache = StreamChangeCache(
            "_group_updates_stream_cache", self._group_updates_id_gen.get_current_token(),
        )

    def __init__(self, db_conn, hs):
        super(SlavedDeviceInboxStore, self).__init__(db_conn, hs)
        self._device_inbox_id_gen = SlavedIdTracker(
            db_conn, "device_max_stream_id", "stream_id",
        )
        self._device_inbox_stream_cache = StreamChangeCache(
            "DeviceInboxStreamChangeCache",
            self._device_inbox_id_gen.get_current_token()
        )
        self._device_federation_outbox_stream_cache = StreamChangeCache(
            "DeviceFederationOutboxStreamChangeCache",
            self._device_inbox_id_gen.get_current_token()
        )

        self._last_device_delete_cache = ExpiringCache(
            cache_name="last_device_delete_cache",
            clock=self._clock,
            max_len=10000,
            expiry_ms=30 * 60 * 1000,
        )

class UserDirectorySlaveStore(
    SlavedEventStore,
    SlavedApplicationServiceStore,
    SlavedRegistrationStore,
    SlavedClientIpStore,
    UserDirectoryStore,
    BaseSlavedStore,
):
    def __init__(self, db_conn, hs):
        super(UserDirectorySlaveStore, self).__init__(db_conn, hs)

        events_max = self._stream_id_gen.get_current_token()
        curr_state_delta_prefill, min_curr_state_delta_id = self._get_cache_dict(
            db_conn, "current_state_delta_stream",
            entity_column="room_id",
            stream_column="stream_id",
            max_value=events_max,  # As we share the stream id with events token
            limit=1000,
        )
        self._curr_state_delta_stream_cache = StreamChangeCache(
            "_curr_state_delta_stream_cache", min_curr_state_delta_id,
            prefilled_cache=curr_state_delta_prefill,
        )

        self._current_state_delta_pos = events_max

    def stream_positions(self):
        result = super(UserDirectorySlaveStore, self).stream_positions()
        result["current_state_deltas"] = self._current_state_delta_pos
        return result

    def process_replication_rows(self, stream_name, token, rows):
        if stream_name == "current_state_deltas":
            self._current_state_delta_pos = token
            for row in rows:
                self._curr_state_delta_stream_cache.entity_has_changed(
                    row.room_id, token
                )
        return super(UserDirectorySlaveStore, self).process_replication_rows(
            stream_name, token, rows
        )

    def test_get_all_entities_changed(self):
        """
        StreamChangeCache.get_all_entities_changed will return all changed
        entities since the given position.  If the position is before the start
        of the known stream, it returns None instead.
        """
        cache = StreamChangeCache("#test", 1)

        cache.entity_has_changed("[email protected]", 2)
        cache.entity_has_changed("[email protected]", 3)
        cache.entity_has_changed("[email protected]", 4)

        self.assertEqual(
            cache.get_all_entities_changed(1),
            ["[email protected]", "[email protected]", "[email protected]"],
        )
        self.assertEqual(
            cache.get_all_entities_changed(2), ["[email protected]", "[email protected]"]
        )
        self.assertEqual(cache.get_all_entities_changed(3), ["[email protected]"])
        self.assertEqual(cache.get_all_entities_changed(0), None)