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": 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 ) 

52 

53 

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

55 """ 

56 Same as above but for a given user 

57 """ 

58 sub_ids = ( 

59 session.execute( 

60 select(PushNotificationSubscription.id) 

61 .where(PushNotificationSubscription.user_id == user_id) 

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

63 ) 

64 .scalars() 

65 .all() 

66 ) 

67 for sub_id in sub_ids: 

68 push_to_subscription( 

69 session, 

70 push_notification_subscription_id=sub_id, 

71 user_id=user_id, 

72 topic_action=topic_action, 

73 key=key, 

74 title=title, 

75 body=body, 

76 icon=icon, 

77 url=url, 

78 ttl=ttl, 

79 ) 

80 

81 

82def push_to_user( 

83 session, 

84 *, 

85 user_id: int, 

86 topic_action: str, 

87 key: str = None, 

88 title: str, 

89 body: str, 

90 icon: str = None, 

91 url: str = None, 

92 ttl: int = 0, 

93): 

94 """ 

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

96 """ 

97 _push_to_user( 

98 session, 

99 user_id=user_id, 

100 topic_action=topic_action, 

101 key=key, 

102 title=title, 

103 body=body, 

104 icon=icon, 

105 url=url, 

106 ttl=ttl, 

107 )