synapse/tests/storage/databases/test_state_store.py
Olivier Wilkinson (reivilibre) 831a7a4592 Add some basic tests about requests getting deduplicated
More to follow when #10825 is ready.
2021-09-21 17:14:19 +01:00

110 lines
3.9 KiB
Python

from typing import Dict, List, Sequence, Tuple
from unittest.mock import patch
from twisted.internet.defer import Deferred, ensureDeferred
from synapse.storage.state import StateFilter
from synapse.types import MutableStateMap, StateMap
from tests.unittest import HomeserverTestCase
class StateGroupInflightCachingTestCase(HomeserverTestCase):
def setUp(self) -> None:
super(StateGroupInflightCachingTestCase, self).setUp()
# Patch out the `_get_state_groups_from_groups`.
# This is useful because it lets us pretend we have a slow database.
gsgfg_patch = patch(
"synapse.storage.databases.state.store.StateGroupDataStore._get_state_groups_from_groups",
self._fake_get_state_groups_from_groups,
)
gsgfg_patch.start()
self.addCleanup(gsgfg_patch.stop)
self.gsgfg_calls: List[
Tuple[Tuple[int, ...], StateFilter, Deferred[Dict[int, StateMap[str]]]]
] = []
def prepare(self, reactor, clock, homeserver) -> None:
super(StateGroupInflightCachingTestCase, self).prepare(
reactor, clock, homeserver
)
self.state_storage = homeserver.get_storage().state
self.state_datastore = homeserver.get_datastores().state
def _fake_get_state_groups_from_groups(
self, groups: Sequence[int], state_filter: StateFilter
) -> "Deferred[Dict[int, StateMap[str]]]":
print("hi", groups, state_filter)
d: Deferred[Dict[int, StateMap[str]]] = Deferred()
self.gsgfg_calls.append((tuple(groups), state_filter, d))
return d
def _complete_request_fake(
self,
groups: Tuple[int, ...],
state_filter: StateFilter,
d: "Deferred[Dict[int, StateMap[str]]]",
) -> None:
"""
Assemble a fake database response and complete the database request.
"""
result: Dict[int, StateMap[str]] = {}
for group in groups:
group_result: MutableStateMap[str] = {}
result[group] = group_result
for state_type, state_keys in state_filter.types.items():
if state_keys is None:
group_result[
(state_type, "wild wombat")
] = f"{group} {state_type} wild wombat"
group_result[
(state_type, "wild spqr")
] = f"{group} {state_type} wild spqr"
else:
for state_key in state_keys:
group_result[
(state_type, state_key)
] = f"{group} {state_type} {state_key}"
if state_filter.include_others:
group_result[("something.else", "wild")] = "card"
d.callback(result)
def test_duplicate_requests_deduplicated(self) -> None:
req1 = ensureDeferred(
self.state_datastore._get_state_for_group_using_inflight_cache(
42, StateFilter.all()
)
)
self.pump(by=0.1)
# This should have gone to the database
self.assertEqual(len(self.gsgfg_calls), 1)
self.assertFalse(req1.called)
req2 = ensureDeferred(
self.state_datastore._get_state_for_group_using_inflight_cache(
42, StateFilter.all()
)
)
self.pump(by=0.1)
# No more calls should have gone to the database
self.assertEqual(len(self.gsgfg_calls), 1)
self.assertFalse(req1.called)
self.assertFalse(req2.called)
groups, sf, d = self.gsgfg_calls[0]
self.assertEqual(groups, (42,))
self.assertEqual(sf, StateFilter.all())
# Now we can complete the request
self._complete_request_fake(groups, sf, d)
self.assertEqual(self.get_success(req1), {("something.else", "wild"): "card"})
self.assertEqual(self.get_success(req2), {("something.else", "wild"): "card"})