1 files changed, 20 insertions, 16 deletions

diff --git a/swiftstory/Client.py b/swiftstory/Client.py
index 33b8e4b..b13e68b 100644
--- a/swiftstory/Client.py
+++ b/swiftstory/Client.py
@@ -1,16 +1,17 @@
+import websockets
+
 from swiftstory.Status import error
 from swiftstory.Game import Game
 
 class Client():
-    def __init__(self, socket, handler, game_manager):
+    def __init__(self, socket, game_manager):
         self.game = None
         self.game_manager = game_manager
 
-        self.handler = handler
         self.socket = socket
         self.player = None
 
-    def join_game(self, game_name, lang):
+    async def join_game(self, game_name, lang):
         if self.game is not None:
             return error('You are already in a game')
 
@@ -23,32 +24,32 @@ class Client():
         if game is None:
             return error('Invalid language')
 
-        return game.try_join(self)
+        return await game.try_join(self)
 
     def set_game(self, game):
         self.game = game
     def set_player(self, player):
         self.player = player
 
-    def play_white_card(self, card_id):
+    async def play_white_card(self, card_id):
         if self.game is None:
             return error('You have to join a game first')
-        return self.game.try_play_card(self.player, card_id)
+        return await self.game.try_play_card(self.player, card_id)
 
-    def pick_black_card(self):
+    async def pick_black_card(self):
         if self.game is None:
             return error('You have to join a game first')
-        return self.game.try_become_judge(self.player)
+        return await self.game.try_become_judge(self.player)
 
-    def collect_cards(self):
+    async def collect_cards(self):
         if self.game is None:
             error('You have to join a game first')
-        return self.game.try_collect_cards(self.player)
+        return await self.game.try_collect_cards(self.player)
 
-    def designate_card(self, card_id):
+    async def designate_card(self, card_id):
         if self.game is None:
             return error('You have to join a game first')
-        return self.game.try_designate_card(self.player, card_id)
+        return await self.game.try_designate_card(self.player, card_id)
 
     def view_player_cards(self):
         if self.game is None:
@@ -65,9 +66,12 @@ class Client():
             return error('You have to join a game first')
         return self.game.try_view_black_card(self.player)
 
-    def send_notification(self, message):
-        self.socket.send_message(self.handler, message)
+    async def send_notification(self, message):
+        try:
+            await self.socket.send(message)
+        except websockets.exceptions.ConnectionClosed:
+            print("Recipient has disconnected.")
 
-    def disconnect(self):
+    async def disconnect(self):
         if self.player is not None:
-            self.game.disconnect(self.player)
+            await self.game.disconnect(self.player)