Coverage for backend/models.py: 97%
100 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-04-17 17:55 +0000
1from __future__ import annotations
3from enum import StrEnum, auto
4from typing import Annotated, Literal
6from pydantic import BaseModel, Field
9class Role(StrEnum):
10 """Player roles."""
12 manager = auto()
13 burger = auto()
14 drink = auto()
15 side = auto()
18class MessageKind(StrEnum):
19 """Message kinds."""
21 initializer = "initializer"
22 game_state = "game_state"
23 lobby_lifecycle = "lobby_lifecycle"
24 chat = "chat"
27class Initializer(BaseModel):
28 """
29 Initializes player state.
31 Links the websocket connection to a specific player and lobby.
33 Attributes:
34 code: Lobby join code.
35 id: Player id.
36 """
38 type: Literal[MessageKind.initializer] = MessageKind.initializer
40 code: tuple[str, ...]
41 id: str
44class LobbyJoinRequest(BaseModel):
45 """
46 Request to join a lobby.
48 Attributes:
49 code: Lobby join code.
50 """
52 code: tuple[str, ...]
55class GameStateUpdateKind(StrEnum):
56 """Game state update kinds."""
58 new_order = "new_order"
59 role_assignment = "role_assignment"
60 order_score = "order_score"
61 order_submission = "order_submission"
62 order_component = "order_component"
63 day_end = "day_end"
66class NewOrder(BaseModel):
67 """A new order for the manager."""
69 type: Literal[MessageKind.game_state] = MessageKind.game_state
70 game_state_update_type: Literal[GameStateUpdateKind.new_order] = GameStateUpdateKind.new_order
72 order: Order
75class RoleAssignment(BaseModel):
76 """Assign a role to the player."""
78 type: Literal[MessageKind.game_state] = MessageKind.game_state
79 game_state_update_type: Literal[GameStateUpdateKind.role_assignment] = GameStateUpdateKind.role_assignment
81 role: Role
84class OrderScore(BaseModel):
85 """A score for the order."""
87 type: Literal[MessageKind.game_state] = MessageKind.game_state
88 game_state_update_type: Literal[GameStateUpdateKind.order_score] = GameStateUpdateKind.order_score
90 score: int
93class OrderComponent(BaseModel):
94 """A finished component is submitted to the manager."""
96 type: Literal[MessageKind.game_state] = MessageKind.game_state
97 game_state_update_type: Literal[GameStateUpdateKind.order_component] = GameStateUpdateKind.order_component
99 component_type: str
100 component: Burger | Side | Drink
103class OrderSubmission(BaseModel):
104 """An order is submitted for review."""
106 type: Literal[MessageKind.game_state] = MessageKind.game_state
107 game_state_update_type: Literal[GameStateUpdateKind.order_submission] = GameStateUpdateKind.order_submission
109 order: Order
112class Order(BaseModel):
113 """A complete order."""
115 burger: Burger | None
116 drink: Drink | None
117 side: Side | None
120class Burger(BaseModel):
121 """A burger."""
123 ingredients: list[str]
126class Drink(BaseModel):
127 """A drink."""
129 color: str
130 fill: float
131 size: str
134class Side(BaseModel):
135 """A side."""
137 table_state: str
140class DayEnd(BaseModel):
141 """No more orders for this day."""
143 type: Literal[MessageKind.game_state] = MessageKind.game_state
144 game_state_update_type: Literal[GameStateUpdateKind.day_end] = GameStateUpdateKind.day_end
146 day: int
149type GameStateUpdate = Annotated[
150 NewOrder | RoleAssignment | OrderScore | OrderSubmission | OrderComponent | DayEnd,
151 Field(discriminator="game_state_update_type"),
152]
155class LobbyLifecycleEventKind(StrEnum):
156 """Lobby lifecycle event kinds."""
158 player_join = "player_join"
159 player_leave = "player_leave"
160 player_count = "player_count"
161 game_start = "game_start"
162 game_end = "game_end"
165class PlayerJoin(BaseModel):
166 """A player joining a lobby."""
168 type: Literal[MessageKind.lobby_lifecycle] = MessageKind.lobby_lifecycle
169 lifecycle_type: Literal[LobbyLifecycleEventKind.player_join] = LobbyLifecycleEventKind.player_join
171 id: str
174class PlayerLeave(BaseModel):
175 """A player leaving a lobby."""
177 type: Literal[MessageKind.lobby_lifecycle] = MessageKind.lobby_lifecycle
178 lifecycle_type: Literal[LobbyLifecycleEventKind.player_leave] = LobbyLifecycleEventKind.player_leave
180 id: str
183class PlayerCount(BaseModel):
184 """Current number of players in the lobby."""
186 type: Literal[MessageKind.lobby_lifecycle] = MessageKind.lobby_lifecycle
187 lifecycle_type: Literal[LobbyLifecycleEventKind.player_count] = LobbyLifecycleEventKind.player_count
189 count: int
190 player_ids: list[str] = Field(default_factory=list)
193class GameStart(BaseModel):
194 """The host starts the game."""
196 type: Literal[MessageKind.lobby_lifecycle] = MessageKind.lobby_lifecycle
197 lifecycle_type: Literal[LobbyLifecycleEventKind.game_start] = LobbyLifecycleEventKind.game_start
200class GameEnd(BaseModel):
201 """All days are completed."""
203 type: Literal[MessageKind.lobby_lifecycle] = MessageKind.lobby_lifecycle
204 lifecycle_type: Literal[LobbyLifecycleEventKind.game_end] = LobbyLifecycleEventKind.game_end
207type LifecycleEvent = Annotated[
208 PlayerJoin | PlayerLeave | PlayerCount | GameStart | GameEnd, Field(discriminator="lifecycle_type")
209]
212class Chat(BaseModel):
213 """The manager is typing."""
215 type: Literal[MessageKind.chat] = MessageKind.chat
217 typing: bool
220class Message(BaseModel):
221 """Wrapper for message models."""
223 data: Initializer | GameStateUpdate | LifecycleEvent | Chat = Field(discriminator="type")
226if __name__ == "__main__":
227 import json
229 message_schema = Message.model_json_schema()
230 print(json.dumps(message_schema, indent=2)) # noqa: T201