From cc42815cced3f3cf1a0ba19d979eb62d5f4d14ff Mon Sep 17 00:00:00 2001
From: legonzaur <legonzaur@live.fr>
Date: Sun, 12 May 2024 11:05:09 +0200
Subject: [PATCH] Attempt to fix #12

---
 src/controller/client.py      | 9 +++++++--
 src/controller/lobbyClient.py | 2 ++
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/controller/client.py b/src/controller/client.py
index 09731af..0aad0c2 100644
--- a/src/controller/client.py
+++ b/src/controller/client.py
@@ -1,6 +1,6 @@
 """Handles communication between one WebSocket tunnel and the game"""
 
-from typing import Callable, Any, TYPE_CHECKING
+from typing import Callable, Any, TYPE_CHECKING, Awaitable
 from pydantic import ValidationError
 
 from websockets.server import WebSocketServerProtocol
@@ -45,6 +45,8 @@ class ClientController:
 
     track_connection: Callable
 
+    on_message = Callable[[Message], Awaitable[None]]
+
     def __init__(
         self,
         game_controller: "GameController",
@@ -53,7 +55,10 @@ class ClientController:
         self.websocket = websocket
         self.game_controller = game_controller
         self.game_controller.clients.append(self)
-        self.on_message = lambda _: None
+        self.on_message = self.on_message_stub
+
+    async def on_message_stub(self, message: Message):
+        pass
 
     async def close(self):
         """Remove self from Game broadcasting list"""
diff --git a/src/controller/lobbyClient.py b/src/controller/lobbyClient.py
index 6b52a87..9f8cd35 100644
--- a/src/controller/lobbyClient.py
+++ b/src/controller/lobbyClient.py
@@ -158,4 +158,6 @@ class LobbyClientController:
                 )
 
     async def game_context(self, message: Message) -> None:
+        if self.game_client is None:
+            return
         await self.game_client.on_message(message)