test_presence_router.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522
  1. # Copyright 2021 The Matrix.org Foundation C.I.C.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the 'License');
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an 'AS IS' BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
  15. from unittest.mock import Mock
  16. import attr
  17. from twisted.test.proto_helpers import MemoryReactor
  18. from synapse.api.constants import EduTypes
  19. from synapse.events.presence_router import PresenceRouter, load_legacy_presence_router
  20. from synapse.federation.units import Transaction
  21. from synapse.handlers.presence import UserPresenceState
  22. from synapse.module_api import ModuleApi
  23. from synapse.rest import admin
  24. from synapse.rest.client import login, presence, room
  25. from synapse.server import HomeServer
  26. from synapse.types import JsonDict, StreamToken, create_requester
  27. from synapse.util import Clock
  28. from tests.handlers.test_sync import generate_sync_config
  29. from tests.test_utils import simple_async_mock
  30. from tests.unittest import (
  31. FederatingHomeserverTestCase,
  32. HomeserverTestCase,
  33. override_config,
  34. )
  35. @attr.s
  36. class PresenceRouterTestConfig:
  37. users_who_should_receive_all_presence = attr.ib(type=List[str], default=[])
  38. class LegacyPresenceRouterTestModule:
  39. def __init__(self, config: PresenceRouterTestConfig, module_api: ModuleApi):
  40. self._config = config
  41. self._module_api = module_api
  42. async def get_users_for_states(
  43. self, state_updates: Iterable[UserPresenceState]
  44. ) -> Dict[str, Set[UserPresenceState]]:
  45. users_to_state = {
  46. user_id: set(state_updates)
  47. for user_id in self._config.users_who_should_receive_all_presence
  48. }
  49. return users_to_state
  50. async def get_interested_users(self, user_id: str) -> Union[Set[str], str]:
  51. if user_id in self._config.users_who_should_receive_all_presence:
  52. return PresenceRouter.ALL_USERS
  53. return set()
  54. @staticmethod
  55. def parse_config(config_dict: dict) -> PresenceRouterTestConfig:
  56. """Parse a configuration dictionary from the homeserver config, do
  57. some validation and return a typed PresenceRouterConfig.
  58. Args:
  59. config_dict: The configuration dictionary.
  60. Returns:
  61. A validated config object.
  62. """
  63. # Initialise a typed config object
  64. config = PresenceRouterTestConfig()
  65. users_who_should_receive_all_presence = config_dict.get(
  66. "users_who_should_receive_all_presence"
  67. )
  68. assert isinstance(users_who_should_receive_all_presence, list)
  69. config.users_who_should_receive_all_presence = (
  70. users_who_should_receive_all_presence
  71. )
  72. return config
  73. class PresenceRouterTestModule:
  74. def __init__(self, config: PresenceRouterTestConfig, api: ModuleApi):
  75. self._config = config
  76. self._module_api = api
  77. api.register_presence_router_callbacks(
  78. get_users_for_states=self.get_users_for_states,
  79. get_interested_users=self.get_interested_users,
  80. )
  81. async def get_users_for_states(
  82. self, state_updates: Iterable[UserPresenceState]
  83. ) -> Dict[str, Set[UserPresenceState]]:
  84. users_to_state = {
  85. user_id: set(state_updates)
  86. for user_id in self._config.users_who_should_receive_all_presence
  87. }
  88. return users_to_state
  89. async def get_interested_users(self, user_id: str) -> Union[Set[str], str]:
  90. if user_id in self._config.users_who_should_receive_all_presence:
  91. return PresenceRouter.ALL_USERS
  92. return set()
  93. @staticmethod
  94. def parse_config(config_dict: dict) -> PresenceRouterTestConfig:
  95. """Parse a configuration dictionary from the homeserver config, do
  96. some validation and return a typed PresenceRouterConfig.
  97. Args:
  98. config_dict: The configuration dictionary.
  99. Returns:
  100. A validated config object.
  101. """
  102. # Initialise a typed config object
  103. config = PresenceRouterTestConfig()
  104. users_who_should_receive_all_presence = config_dict.get(
  105. "users_who_should_receive_all_presence"
  106. )
  107. assert isinstance(users_who_should_receive_all_presence, list)
  108. config.users_who_should_receive_all_presence = (
  109. users_who_should_receive_all_presence
  110. )
  111. return config
  112. class PresenceRouterTestCase(FederatingHomeserverTestCase):
  113. """
  114. Test cases using a custom PresenceRouter
  115. By default in test cases, federation sending is disabled. This class re-enables it
  116. for the main process by setting `federation_sender_instances` to None.
  117. """
  118. servlets = [
  119. admin.register_servlets,
  120. login.register_servlets,
  121. room.register_servlets,
  122. presence.register_servlets,
  123. ]
  124. def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
  125. # Mock out the calls over federation.
  126. self.fed_transport_client = Mock(spec=["send_transaction"])
  127. self.fed_transport_client.send_transaction = simple_async_mock({})
  128. hs = self.setup_test_homeserver(
  129. federation_transport_client=self.fed_transport_client,
  130. )
  131. load_legacy_presence_router(hs)
  132. return hs
  133. def prepare(
  134. self, reactor: MemoryReactor, clock: Clock, homeserver: HomeServer
  135. ) -> None:
  136. self.sync_handler = self.hs.get_sync_handler()
  137. self.module_api = homeserver.get_module_api()
  138. def default_config(self) -> JsonDict:
  139. config = super().default_config()
  140. config["federation_sender_instances"] = None
  141. return config
  142. @override_config(
  143. {
  144. "presence": {
  145. "presence_router": {
  146. "module": __name__ + ".LegacyPresenceRouterTestModule",
  147. "config": {
  148. "users_who_should_receive_all_presence": [
  149. "@presence_gobbler:test",
  150. ]
  151. },
  152. }
  153. },
  154. }
  155. )
  156. def test_receiving_all_presence_legacy(self) -> None:
  157. self.receiving_all_presence_test_body()
  158. @override_config(
  159. {
  160. "modules": [
  161. {
  162. "module": __name__ + ".PresenceRouterTestModule",
  163. "config": {
  164. "users_who_should_receive_all_presence": [
  165. "@presence_gobbler:test",
  166. ]
  167. },
  168. },
  169. ],
  170. }
  171. )
  172. def test_receiving_all_presence(self) -> None:
  173. self.receiving_all_presence_test_body()
  174. def receiving_all_presence_test_body(self) -> None:
  175. """Test that a user that does not share a room with another other can receive
  176. presence for them, due to presence routing.
  177. """
  178. # Create a user who should receive all presence of others
  179. self.presence_receiving_user_id = self.register_user(
  180. "presence_gobbler", "monkey"
  181. )
  182. self.presence_receiving_user_tok = self.login("presence_gobbler", "monkey")
  183. # And two users who should not have any special routing
  184. self.other_user_one_id = self.register_user("other_user_one", "monkey")
  185. self.other_user_one_tok = self.login("other_user_one", "monkey")
  186. self.other_user_two_id = self.register_user("other_user_two", "monkey")
  187. self.other_user_two_tok = self.login("other_user_two", "monkey")
  188. # Put the other two users in a room with each other
  189. room_id = self.helper.create_room_as(
  190. self.other_user_one_id, tok=self.other_user_one_tok
  191. )
  192. self.helper.invite(
  193. room_id,
  194. self.other_user_one_id,
  195. self.other_user_two_id,
  196. tok=self.other_user_one_tok,
  197. )
  198. self.helper.join(room_id, self.other_user_two_id, tok=self.other_user_two_tok)
  199. # User one sends some presence
  200. send_presence_update(
  201. self,
  202. self.other_user_one_id,
  203. self.other_user_one_tok,
  204. "online",
  205. "boop",
  206. )
  207. # Check that the presence receiving user gets user one's presence when syncing
  208. presence_updates, sync_token = sync_presence(
  209. self, self.presence_receiving_user_id
  210. )
  211. self.assertEqual(len(presence_updates), 1)
  212. presence_update: UserPresenceState = presence_updates[0]
  213. self.assertEqual(presence_update.user_id, self.other_user_one_id)
  214. self.assertEqual(presence_update.state, "online")
  215. self.assertEqual(presence_update.status_msg, "boop")
  216. # Have all three users send presence
  217. send_presence_update(
  218. self,
  219. self.other_user_one_id,
  220. self.other_user_one_tok,
  221. "online",
  222. "user_one",
  223. )
  224. send_presence_update(
  225. self,
  226. self.other_user_two_id,
  227. self.other_user_two_tok,
  228. "online",
  229. "user_two",
  230. )
  231. send_presence_update(
  232. self,
  233. self.presence_receiving_user_id,
  234. self.presence_receiving_user_tok,
  235. "online",
  236. "presence_gobbler",
  237. )
  238. # Check that the presence receiving user gets everyone's presence
  239. presence_updates, _ = sync_presence(
  240. self, self.presence_receiving_user_id, sync_token
  241. )
  242. self.assertEqual(len(presence_updates), 3)
  243. # But that User One only get itself and User Two's presence
  244. presence_updates, _ = sync_presence(self, self.other_user_one_id)
  245. self.assertEqual(len(presence_updates), 2)
  246. found = False
  247. for update in presence_updates:
  248. if update.user_id == self.other_user_two_id:
  249. self.assertEqual(update.state, "online")
  250. self.assertEqual(update.status_msg, "user_two")
  251. found = True
  252. self.assertTrue(found)
  253. @override_config(
  254. {
  255. "presence": {
  256. "presence_router": {
  257. "module": __name__ + ".LegacyPresenceRouterTestModule",
  258. "config": {
  259. "users_who_should_receive_all_presence": [
  260. "@presence_gobbler1:test",
  261. "@presence_gobbler2:test",
  262. "@far_away_person:island",
  263. ]
  264. },
  265. }
  266. },
  267. }
  268. )
  269. def test_send_local_online_presence_to_with_module_legacy(self) -> None:
  270. self.send_local_online_presence_to_with_module_test_body()
  271. @override_config(
  272. {
  273. "modules": [
  274. {
  275. "module": __name__ + ".PresenceRouterTestModule",
  276. "config": {
  277. "users_who_should_receive_all_presence": [
  278. "@presence_gobbler1:test",
  279. "@presence_gobbler2:test",
  280. "@far_away_person:island",
  281. ]
  282. },
  283. },
  284. ],
  285. }
  286. )
  287. def test_send_local_online_presence_to_with_module(self) -> None:
  288. self.send_local_online_presence_to_with_module_test_body()
  289. def send_local_online_presence_to_with_module_test_body(self) -> None:
  290. """Tests that send_local_presence_to_users sends local online presence to a set
  291. of specified local and remote users, with a custom PresenceRouter module enabled.
  292. """
  293. # Create a user who will send presence updates
  294. self.other_user_id = self.register_user("other_user", "monkey")
  295. self.other_user_tok = self.login("other_user", "monkey")
  296. # And another two users that will also send out presence updates, as well as receive
  297. # theirs and everyone else's
  298. self.presence_receiving_user_one_id = self.register_user(
  299. "presence_gobbler1", "monkey"
  300. )
  301. self.presence_receiving_user_one_tok = self.login("presence_gobbler1", "monkey")
  302. self.presence_receiving_user_two_id = self.register_user(
  303. "presence_gobbler2", "monkey"
  304. )
  305. self.presence_receiving_user_two_tok = self.login("presence_gobbler2", "monkey")
  306. # Have all three users send some presence updates
  307. send_presence_update(
  308. self,
  309. self.other_user_id,
  310. self.other_user_tok,
  311. "online",
  312. "I'm online!",
  313. )
  314. send_presence_update(
  315. self,
  316. self.presence_receiving_user_one_id,
  317. self.presence_receiving_user_one_tok,
  318. "online",
  319. "I'm also online!",
  320. )
  321. send_presence_update(
  322. self,
  323. self.presence_receiving_user_two_id,
  324. self.presence_receiving_user_two_tok,
  325. "unavailable",
  326. "I'm in a meeting!",
  327. )
  328. # Mark each presence-receiving user for receiving all user presence
  329. self.get_success(
  330. self.module_api.send_local_online_presence_to(
  331. [
  332. self.presence_receiving_user_one_id,
  333. self.presence_receiving_user_two_id,
  334. ]
  335. )
  336. )
  337. # Perform a sync for each user
  338. # The other user should only receive their own presence
  339. presence_updates, _ = sync_presence(self, self.other_user_id)
  340. self.assertEqual(len(presence_updates), 1)
  341. presence_update: UserPresenceState = presence_updates[0]
  342. self.assertEqual(presence_update.user_id, self.other_user_id)
  343. self.assertEqual(presence_update.state, "online")
  344. self.assertEqual(presence_update.status_msg, "I'm online!")
  345. # Whereas both presence receiving users should receive everyone's presence updates
  346. presence_updates, _ = sync_presence(self, self.presence_receiving_user_one_id)
  347. self.assertEqual(len(presence_updates), 3)
  348. presence_updates, _ = sync_presence(self, self.presence_receiving_user_two_id)
  349. self.assertEqual(len(presence_updates), 3)
  350. # We stagger sending of presence, so we need to wait a bit for them to
  351. # get sent out.
  352. self.reactor.advance(60)
  353. # Test that sending to a remote user works
  354. remote_user_id = "@far_away_person:island"
  355. # Note that due to the remote user being in our module's
  356. # users_who_should_receive_all_presence config, they would have
  357. # received user presence updates already.
  358. #
  359. # Thus we reset the mock, and try sending all online local user
  360. # presence again
  361. self.fed_transport_client.send_transaction.reset_mock()
  362. # Broadcast local user online presence
  363. self.get_success(
  364. self.module_api.send_local_online_presence_to([remote_user_id])
  365. )
  366. # We stagger sending of presence, so we need to wait a bit for them to
  367. # get sent out.
  368. self.reactor.advance(60)
  369. # Check that the expected presence updates were sent
  370. # We explicitly compare using sets as we expect that calling
  371. # module_api.send_local_online_presence_to will create a presence
  372. # update that is a duplicate of the specified user's current presence.
  373. # These are sent to clients and will be picked up below, thus we use a
  374. # set to deduplicate. We're just interested that non-offline updates were
  375. # sent out for each user ID.
  376. expected_users = {
  377. self.other_user_id,
  378. self.presence_receiving_user_one_id,
  379. self.presence_receiving_user_two_id,
  380. }
  381. found_users = set()
  382. calls = self.fed_transport_client.send_transaction.call_args_list
  383. for call in calls:
  384. call_args = call[0]
  385. federation_transaction: Transaction = call_args[0]
  386. # Get the sent EDUs in this transaction
  387. edus = federation_transaction.get_dict()["edus"]
  388. for edu in edus:
  389. # Make sure we're only checking presence-type EDUs
  390. if edu["edu_type"] != EduTypes.PRESENCE:
  391. continue
  392. # EDUs can contain multiple presence updates
  393. for presence_edu in edu["content"]["push"]:
  394. # Check for presence updates that contain the user IDs we're after
  395. found_users.add(presence_edu["user_id"])
  396. # Ensure that no offline states are being sent out
  397. self.assertNotEqual(presence_edu["presence"], "offline")
  398. self.assertEqual(found_users, expected_users)
  399. def send_presence_update(
  400. testcase: HomeserverTestCase,
  401. user_id: str,
  402. access_token: str,
  403. presence_state: str,
  404. status_message: Optional[str] = None,
  405. ) -> JsonDict:
  406. # Build the presence body
  407. body = {"presence": presence_state}
  408. if status_message:
  409. body["status_msg"] = status_message
  410. # Update the user's presence state
  411. channel = testcase.make_request(
  412. "PUT", "/presence/%s/status" % (user_id,), body, access_token=access_token
  413. )
  414. testcase.assertEqual(channel.code, 200)
  415. return channel.json_body
  416. def sync_presence(
  417. testcase: HomeserverTestCase,
  418. user_id: str,
  419. since_token: Optional[StreamToken] = None,
  420. ) -> Tuple[List[UserPresenceState], StreamToken]:
  421. """Perform a sync request for the given user and return the user presence updates
  422. they've received, as well as the next_batch token.
  423. This method assumes testcase.sync_handler points to the homeserver's sync handler.
  424. Args:
  425. testcase: The testcase that is currently being run.
  426. user_id: The ID of the user to generate a sync response for.
  427. since_token: An optional token to indicate from at what point to sync from.
  428. Returns:
  429. A tuple containing a list of presence updates, and the sync response's
  430. next_batch token.
  431. """
  432. requester = create_requester(user_id)
  433. sync_config = generate_sync_config(requester.user.to_string())
  434. sync_result = testcase.get_success(
  435. testcase.hs.get_sync_handler().wait_for_sync_for_user(
  436. requester, sync_config, since_token
  437. )
  438. )
  439. return sync_result.presence, sync_result.next_batch