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

1import functools 

2import json 

3 

4from sqlalchemy.orm import Session 

5from sqlalchemy.sql import func 

6 

7from couchers import urls 

8from couchers.config import config 

9from couchers.jobs.enqueue import queue_job 

10from couchers.models import PushNotificationSubscription 

11from couchers.notifications.push_api import get_vapid_public_key_from_private_key 

12from couchers.proto.internal import jobs_pb2 

13from couchers.sql import couchers_select as select 

14 

15 

16@functools.cache 

17def get_vapid_public_key() -> str: 

18 return get_vapid_public_key_from_private_key(config["PUSH_NOTIFICATIONS_VAPID_PRIVATE_KEY"]) 

19 

20 

21def push_to_subscription( 

22 session: Session, 

23 *, 

24 push_notification_subscription_id: int, 

25 user_id: int, 

26 topic_action: str, 

27 key: str | None = None, 

28 title: str = "", 

29 body: str = "", 

30 icon: str | None = None, 

31 url: str | None = None, 

32 ttl: int = 0, 

33) -> None: 

34 queue_job( 

35 session, 

36 job_type="send_raw_push_notification", 

37 payload=jobs_pb2.SendRawPushNotificationPayload( 

38 data=json.dumps( 

39 { 

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

41 "body": body[:2000], 

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

43 "url": url, 

44 "user_id": user_id, 

45 "topic_action": topic_action, 

46 "key": key or "", 

47 } 

48 ).encode("utf8"), 

49 push_notification_subscription_id=push_notification_subscription_id, 

50 ttl=ttl, 

51 ), 

52 priority=7, 

53 ) 

54 

55 

56def _push_to_user( 

57 session: Session, 

58 user_id: int, 

59 topic_action: str, 

60 key: str | None, 

61 title: str, 

62 body: str, 

63 icon: str | None, 

64 url: str | None, 

65 ttl: int, 

66) -> None: 

67 """ 

68 Same as above but for a given user 

69 """ 

70 sub_ids = ( 

71 session.execute( 

72 select(PushNotificationSubscription.id) 

73 .where(PushNotificationSubscription.user_id == user_id) 

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

75 ) 

76 .scalars() 

77 .all() 

78 ) 

79 for sub_id in sub_ids: 

80 push_to_subscription( 

81 session, 

82 push_notification_subscription_id=sub_id, 

83 user_id=user_id, 

84 topic_action=topic_action, 

85 key=key, 

86 title=title, 

87 body=body, 

88 icon=icon, 

89 url=url, 

90 ttl=ttl, 

91 ) 

92 

93 

94def push_to_user( 

95 session: Session, 

96 *, 

97 user_id: int, 

98 topic_action: str, 

99 key: str | None = None, 

100 title: str = "", 

101 body: str = "", 

102 icon: str | None = None, 

103 url: str | None = None, 

104 ttl: int = 0, 

105) -> None: 

106 """ 

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

108 """ 

109 _push_to_user( 

110 session, 

111 user_id=user_id, 

112 topic_action=topic_action, 

113 key=key, 

114 title=title, 

115 body=body, 

116 icon=icon, 

117 url=url, 

118 ttl=ttl, 

119 )