Coverage for microservice_websocket/app/services/database/models.py: 91%

1from __future__ import annotations 

2 

3from datetime import datetime 

4 

5from app.utils.exceptions import NotFoundException 

6from beanie import Document, PydanticObjectId 

7from beanie.exceptions import DocumentWasNotSaved 

8from beanie.operators import And, Eq 

9from pydantic import BaseModel, Field 

10 

11from ...models.node_settings import DetectorSettings 

12from ...utils.enums import NodeState 

13 

14 

15class CustomDocument(Document): 

16 @property 

17 def id(self) -> PydanticObjectId: 

18 obj_id = super().id 

19 if obj_id is None: 

20 raise DocumentWasNotSaved 

21 

22 return obj_id 

23 

24 

25class Organization(CustomDocument): 

26 organizationName: str 

27 

28 class Serialized(BaseModel): 

29 id: str 

30 organizationName: str 

31 

32 def serialize(self) -> Organization.Serialized: 

33 return Organization.Serialized( 

34 id=str(self.id), organizationName=self.organizationName 

35 ) 

36 

37 

38class Application(CustomDocument): 

39 applicationName: str 

40 organization: PydanticObjectId 

41 

42 class Serialized(BaseModel): 

43 id: str 

44 applicationName: str 

45 organization: str 

46 

47 def serialize(self) -> Application.Serialized: 

48 return Application.Serialized( 

49 id=str(self.id), 

50 applicationName=self.applicationName, 

51 organization=str(self.organization), 

52 ) 

53 

54 

55class User(CustomDocument): 

56 email: str 

57 hashed_password: str 

58 first_name: str 

59 last_name: str 

60 role: str = "standard" 

61 

62 class Serialized(BaseModel): 

63 id: str 

64 email: str 

65 first_name: str 

66 last_name: str 

67 role: str 

68 

69 def serialize(self) -> User.Serialized: 

70 return User.Serialized( 

71 id=str(self.id), 

72 email=self.email, 

73 first_name=self.first_name, 

74 last_name=self.last_name, 

75 role=self.role, 

76 ) 

77 

78 

79class NodeSettings(CustomDocument): 

80 node: PydanticObjectId 

81 d1: DetectorSettings | None = None 

82 d2: DetectorSettings | None = None 

83 d3: DetectorSettings | None = None 

84 d4: DetectorSettings | None = None 

85 

86 class Serialized(BaseModel): 

87 d1: DetectorSettings | None 

88 d2: DetectorSettings | None 

89 d3: DetectorSettings | None 

90 d4: DetectorSettings | None 

91 

92 def serialize(self) -> NodeSettings.Serialized: 

93 return NodeSettings.Serialized( 

94 d1=self.d1, 

95 d2=self.d2, 

96 d3=self.d3, 

97 d4=self.d4, 

98 ) 

99 

100 

101class Node(CustomDocument): 

102 nodeID: int 

103 nodeName: str 

104 application: PydanticObjectId 

105 state: NodeState 

106 lastSeenAt: datetime 

107 

108 class Serialized(BaseModel): 

109 nodeID: int 

110 nodeName: str 

111 application: str 

112 state: str 

113 lastSeenAt: int 

114 unhandledAlertIDs: list[str] 

115 

116 async def serialize(self) -> Node.Serialized: 

117 unhandledAlerts = await Alert.find( 

118 And(Eq(Alert.node, self.id), Eq(Alert.isHandled, False)) 

119 ).to_list() 

120 unhandledAlertIDs = [str(x.id) for x in unhandledAlerts] 

121 

122 return Node.Serialized( 

123 nodeID=self.nodeID, 

124 nodeName=self.nodeName, 

125 application=str(self.application), 

126 state=NodeState.to_irma_ui_state(self.state), 

127 lastSeenAt=int(self.lastSeenAt.timestamp()), 

128 unhandledAlertIDs=unhandledAlertIDs, 

129 ) 

130 

131 

132class Reading(CustomDocument): 

133 node: PydanticObjectId 

134 canID: int = Field(default=..., lt=5, gt=0) 

135 sensorNumber: int = Field(default=..., lt=3, gt=0) 

136 readingID: int 

137 sessionID: int 

138 dangerLevel: int = 0 

139 window1: int = 0 

140 window2: int = 0 

141 window3: int = 0 

142 publishedAt: datetime 

143 

144 class Serialized(BaseModel): 

145 nodeID: int 

146 canID: int 

147 sensorNumber: int 

148 readingID: int 

149 sessionID: int 

150 dangerLevel: int 

151 window1: int 

152 window2: int 

153 window3: int 

154 publishedAt: int 

155 

156 async def serialize(self) -> Reading.Serialized: 

157 node = await Node.get(self.node) 

158 if node is None: 158 ↛ 159line 158 didn't jump to line 159, because the condition on line 158 was never true

159 raise NotFoundException("Node") 

160 

161 return Reading.Serialized( 

162 nodeID=node.nodeID, 

163 canID=self.canID, 

164 sensorNumber=self.sensorNumber, 

165 readingID=self.readingID, 

166 sessionID=self.sessionID, 

167 dangerLevel=self.dangerLevel, 

168 window1=self.window1, 

169 window2=self.window2, 

170 window3=self.window3, 

171 publishedAt=int(self.publishedAt.timestamp()), 

172 ) 

173 

174 

175class Alert(CustomDocument): 

176 reading: PydanticObjectId 

177 node: PydanticObjectId 

178 sessionID: int 

179 isHandled: bool = False 

180 raisedAt: datetime 

181 isConfirmed: bool = False 

182 handledBy: PydanticObjectId | None = None 

183 handledAt: datetime | None = None 

184 handleNote: str = "" 

185 

186 class Serialized(BaseModel): 

187 id: str 

188 nodeID: int 

189 sessionID: int 

190 readingID: int 

191 canID: int 

192 raisedAt: int 

193 

194 async def serialize(self) -> Alert.Serialized: 

195 node = await Node.get(self.node) 

196 if node is None: 196 ↛ 197line 196 didn't jump to line 197, because the condition on line 196 was never true

197 raise NotFoundException("Node") 

198 

199 reading = await Reading.get(self.reading) 

200 if reading is None: 200 ↛ 201line 200 didn't jump to line 201, because the condition on line 200 was never true

201 raise NotFoundException("Reading") 

202 

203 return Alert.Serialized( 

204 id=str(self.id), 

205 nodeID=node.nodeID, 

206 sessionID=self.sessionID, 

207 readingID=reading.readingID, 

208 canID=reading.canID, 

209 raisedAt=int(self.raisedAt.timestamp()), 

210 )