Coverage for backend/models.py: 97%
102 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-05-02 01:42 +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
147 customers_served: int
148 score: int
151type GameStateUpdate = Annotated[
152 NewOrder | RoleAssignment | OrderScore | OrderSubmission | OrderComponent | DayEnd,
153 Field(discriminator="game_state_update_type"),
154]
157class LobbyLifecycleEventKind(StrEnum):
158 """Lobby lifecycle event kinds."""
160 player_join = "player_join"
161 player_leave = "player_leave"
162 player_count = "player_count"
163 game_start = "game_start"
164 game_end = "game_end"
167class PlayerJoin(BaseModel):
168 """A player joining a lobby."""
170 type: Literal[MessageKind.lobby_lifecycle] = MessageKind.lobby_lifecycle
171 lifecycle_type: Literal[LobbyLifecycleEventKind.player_join] = LobbyLifecycleEventKind.player_join
173 id: str
176class PlayerLeave(BaseModel):
177 """A player leaving a lobby."""
179 type: Literal[MessageKind.lobby_lifecycle] = MessageKind.lobby_lifecycle
180 lifecycle_type: Literal[LobbyLifecycleEventKind.player_leave] = LobbyLifecycleEventKind.player_leave
182 id: str
185class PlayerCount(BaseModel):
186 """Current number of players in the lobby."""
188 type: Literal[MessageKind.lobby_lifecycle] = MessageKind.lobby_lifecycle
189 lifecycle_type: Literal[LobbyLifecycleEventKind.player_count] = LobbyLifecycleEventKind.player_count
191 count: int
192 player_ids: list[str] = Field(default_factory=list)
195class GameStart(BaseModel):
196 """The host starts the game."""
198 type: Literal[MessageKind.lobby_lifecycle] = MessageKind.lobby_lifecycle
199 lifecycle_type: Literal[LobbyLifecycleEventKind.game_start] = LobbyLifecycleEventKind.game_start
202class GameEnd(BaseModel):
203 """All days are completed."""
205 type: Literal[MessageKind.lobby_lifecycle] = MessageKind.lobby_lifecycle
206 lifecycle_type: Literal[LobbyLifecycleEventKind.game_end] = LobbyLifecycleEventKind.game_end
209type LifecycleEvent = Annotated[
210 PlayerJoin | PlayerLeave | PlayerCount | GameStart | GameEnd, Field(discriminator="lifecycle_type")
211]
214class Chat(BaseModel):
215 """The manager is typing."""
217 type: Literal[MessageKind.chat] = MessageKind.chat
219 typing: bool
222class Message(BaseModel):
223 """Wrapper for message models."""
225 data: Initializer | GameStateUpdate | LifecycleEvent | Chat = Field(discriminator="type")
228if __name__ == "__main__":
229 import json
231 message_schema = Message.model_json_schema()
232 print(json.dumps(message_schema, indent=2)) # noqa: T201