Coverage for src/couchers/helpers/badges.py: 100%

1from sqlalchemy.orm import Session 

2from sqlalchemy.sql import delete 

3 

4from couchers.context import make_background_user_context 

5from couchers.models import UserBadge 

6from couchers.notifications.notify import notify 

7from couchers.proto import notification_data_pb2 

8from couchers.resources import get_badge_dict 

9 

10 

11def user_add_badge(session: Session, user_id: int, badge_id: str, do_notify: bool = True) -> None: 

12 badge = get_badge_dict()[badge_id] 

13 session.add(UserBadge(user_id=user_id, badge_id=badge_id)) 

14 session.flush() 

15 if do_notify: 

16 context = make_background_user_context(user_id=user_id) 

17 notify( 

18 session, 

19 user_id=user_id, 

20 topic_action="badge:add", 

21 key=badge_id, 

22 data=notification_data_pb2.BadgeAdd( 

23 badge_id=badge["id"], 

24 badge_name=context.get_localized_string("badges", f"{badge_id}_name"), 

25 badge_description=context.get_localized_string("badges", f"{badge_id}_description"), 

26 ), 

27 ) 

28 session.commit() 

29 

30 

31def user_remove_badge(session: Session, user_id: int, badge_id: str) -> None: 

32 badge = get_badge_dict()[badge_id] 

33 session.execute(delete(UserBadge).where(UserBadge.user_id == user_id, UserBadge.badge_id == badge_id)) 

34 session.flush() 

35 context = make_background_user_context(user_id=user_id) 

36 notify( 

37 session, 

38 user_id=user_id, 

39 topic_action="badge:remove", 

40 key=badge_id, 

41 data=notification_data_pb2.BadgeRemove( 

42 badge_id=badge["id"], 

43 badge_name=context.get_localized_string("badges", f"{badge_id}_name"), 

44 badge_description=context.get_localized_string("badges", f"{badge_id}_description"), 

45 ), 

46 ) 

47 session.commit()