Coverage for src/couchers/notifications/push.py: 95%

1import functools 

2import json 

3 

4from sqlalchemy.sql import func 

5 

6from couchers import urls 

7from couchers.config import config 

8from couchers.jobs.enqueue import queue_job 

9from couchers.models import PushNotificationSubscription 

10from couchers.notifications.push_api import get_vapid_public_key_from_private_key 

11from couchers.sql import couchers_select as select 

12from proto.internal import jobs_pb2 

13 

14 

15@functools.cache 

16def get_vapid_public_key(): 

17 return get_vapid_public_key_from_private_key(config["PUSH_NOTIFICATIONS_VAPID_PRIVATE_KEY"]) 

18 

19 

20def push_to_subscription( 

21 session, 

22 *, 

23 push_notification_subscription_id: int, 

24 user_id: int, 

25 topic_action: str, 

26 key: str = None, 

27 title: str, 

28 body: str, 

29 icon: str = None, 

30 url: str = None, 

31 ttl: int = 0, 

32): 

33 queue_job( 

34 session, 

35 job_type="send_raw_push_notification", 

36 payload=jobs_pb2.SendRawPushNotificationPayload( 

37 data=json.dumps( 

38 { 

39 "title": config["NOTIFICATION_PREFIX"] + title[:500], 

40 "body": body[:2000], 

41 "icon": icon or urls.icon_url(), 

42 "url": url, 

43 "user_id": user_id, 

44 "topic_action": topic_action, 

45 "key": key or "", 

46 } 

47 ).encode("utf8"), 

48 push_notification_subscription_id=push_notification_subscription_id, 

49 ttl=ttl, 

50 ), 

51 priority=7, 

52 ) 

53 

54 

55def _push_to_user(session, user_id, topic_action, key, title, body, icon, url, ttl): 

56 """ 

57 Same as above but for a given user 

58 """ 

59 sub_ids = ( 

60 session.execute( 

61 select(PushNotificationSubscription.id) 

62 .where(PushNotificationSubscription.user_id == user_id) 

63 .where(PushNotificationSubscription.disabled_at > func.now()) 

64 ) 

65 .scalars() 

66 .all() 

67 ) 

68 for sub_id in sub_ids: 

69 push_to_subscription( 

70 session, 

71 push_notification_subscription_id=sub_id, 

72 user_id=user_id, 

73 topic_action=topic_action, 

74 key=key, 

75 title=title, 

76 body=body, 

77 icon=icon, 

78 url=url, 

79 ttl=ttl, 

80 ) 

81 

82 

83def push_to_user( 

84 session, 

85 *, 

86 user_id: int, 

87 topic_action: str, 

88 key: str = None, 

89 title: str, 

90 body: str, 

91 icon: str = None, 

92 url: str = None, 

93 ttl: int = 0, 

94): 

95 """ 

96 This indirection is so that this can be easily mocked. Not sure how to do it better :( 

97 """ 

98 _push_to_user( 

99 session, 

100 user_id=user_id, 

101 topic_action=topic_action, 

102 key=key, 

103 title=title, 

104 body=body, 

105 icon=icon, 

106 url=url, 

107 ttl=ttl, 

108 )