Coverage for backend/django/notifications/consumers/NotificationsConsumer.py: 84%

1"""Websocket consumer that streams flowsheet notifications to authenticated users.""" 

2 

3from channels.generic.websocket import AsyncWebsocketConsumer 

4from django.core.exceptions import ObjectDoesNotExist 

5from CoreRoot import settings 

6from CoreRoot.helpers import get_asgi_header_value 

7from authentication.user.AccessTable import AccessTable 

8from authentication.user.models import User 

9from core.auxiliary.models.Flowsheet import Flowsheet 

10 

11async def _get_user(username): 

12 """Return the minimal user record needed for websocket authentication.""" 

13 user = await User.objects.only("id").aget(username=username) 

14 

15 return user 

16 

17 

18async def _get_project_access_entry_for_user(user_id: int, project_id: int) -> AccessTable | None: 

19 return await ( 

20 AccessTable.objects 

21 .filter(project_id=project_id, user_id=user_id) 

22 .afirst() 

23 ) 

24 

25 

26class NotificationsConsumer(AsyncWebsocketConsumer): 

27 """Manage a flowsheet-scoped websocket connection for notification delivery.""" 

28 

29 def __init__(self, *args, **kwargs): 

30 super().__init__(*args, **kwargs) 

31 

32 self.user_id = None 

33 self.flowsheet_id = None 

34 

35 async def connect(self): 

36 """Authenticate the user and join the flowsheet broadcast group.""" 

37 

38 username = get_asgi_header_value(self.scope["headers"], settings.ASGI_REMOTE_USER_HEADER) 

39 raw_flowsheet_id = self.scope["query_params"].get("flowsheetId", [None])[0] 

40 

41 if not username: 41 ↛ 42line 41 didn't jump to line 42 because the condition on line 41 was never true

42 await self.close(reason="Unauthorized") 

43 return 

44 

45 if not raw_flowsheet_id: 45 ↛ 46line 45 didn't jump to line 46 because the condition on line 45 was never true

46 await self.close(reason="Missing flowsheet_id query parameter") 

47 return 

48 

49 try: 

50 self.flowsheet_id = int(raw_flowsheet_id) 

51 except (TypeError, ValueError): 

52 await self.close(reason="Invalid flowsheet_id query parameter") 

53 return 

54 

55 try: 

56 user = await _get_user(username) 

57 flowsheet = await Flowsheet.objects.only("id", "project_id").aget(id=self.flowsheet_id) 

58 except ObjectDoesNotExist: 

59 await self.close(reason="Unauthorized") 

60 return 

61 

62 if flowsheet.project_id is None: 62 ↛ 63line 62 didn't jump to line 63 because the condition on line 62 was never true

63 await self.close(reason="Unauthorized") 

64 return 

65 

66 access_entry = await _get_project_access_entry_for_user(user.id, flowsheet.project_id) 

67 self.user_id = user.id 

68 

69 if access_entry is None: 

70 await self.close(reason="Unauthorized") 

71 return 

72 

73 await self.accept() 

74 

75 # Register this socket with the flowsheet-specific broadcast group so 

76 # `broadcast_view` can fan out to every active session for the user. 

77 await self.channel_layer.group_add(f"{self.flowsheet_id}", self.channel_name) 

78 

79 async def receive(self, text_data=None, bytes_data=None): 

80 """Respond to heartbeat pings from the client to keep the socket alive.""" 

81 if text_data == "ping": 81 ↛ exitline 81 didn't return from function 'receive' because the condition on line 81 was always true

82 await self.send(text_data="pong") 

83 

84 async def disconnect(self, code): 

85 """Remove the socket from the flowsheet broadcast group on close.""" 

86 if self.user_id and self.flowsheet_id: 86 ↛ exitline 86 didn't return from function 'disconnect' because the condition on line 86 was always true

87 await self.channel_layer.group_discard(f"{self.flowsheet_id}", self.channel_name) 

88 

89 async def flowsheet_message(self, event): 

90 """Forward a channel-layer broadcast event payload to the client.""" 

91 await self.send(text_data=event["data"])