test_presence_router.py 18 KB

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