Add timeout to Controller's connect()

yaboli/connection.py

@@ -77,6 +77,7 @@ class Connection:
 			except:
 				pass # errors are not useful here
 			
+			self._pingtask.cancel()
 			await self._pingtask # should stop now that the ws is closed
 			self._ws = None
 	
@@ -91,14 +92,13 @@ class Connection:
 				wait_for_reply = await self._ws.ping()
 				await asyncio.wait_for(wait_for_reply, self.ping_timeout)
 				logger.debug("Pinged!")
+				await asyncio.sleep(self.ping_delay)
 			except asyncio.TimeoutError:
 				logger.warning("Ping timed out.")
 				await self._ws.close()
 				break
-			except (websockets.ConnectionClosed, ConnectionResetError):
-				break
-			else:
-				await asyncio.sleep(self.ping_delay)
+			except (websockets.ConnectionClosed, ConnectionResetError, asyncio.CancelledError):
+				return
 	
 	async def stop(self):
 		"""

yaboli/controller.py

@@ -27,11 +27,12 @@ class Controller:
 	
 	"""
 	
-	def __init__(self, nick, human=False, cookie=None):
+	def __init__(self, nick, human=False, cookie=None, connect_timeout=10):
 		"""
 		roomname        - name of room to connect to
 		human           - whether the human flag should be set on connections
 		cookie          - cookie to use in HTTP request, if any
+		connect_timeout - time for authentication to complete
 		"""
 		self.nick = nick
 		self.human = human
@@ -41,6 +42,7 @@ class Controller:
 		self.password = None
 		
 		self.room = None
+		self.connect_timeout = connect_timeout # in seconds
 		self._connect_result = None
 	
 	def _create_room(self, roomname):
@@ -70,6 +72,7 @@ class Controller:
 		  "no password"    = password needed to connect to room
 		  "wrong password" = password given does not work
 		  "disconnected"   = connection closed before client could access the room
+		  "timeout"        = timed out while waiting for server
 		  "success"        = no failure
 		"""
 		
@@ -97,9 +100,13 @@ class Controller:
 		
 		# connection succeeded, now we need to know whether we can log in
 		# wait for success/authentication/disconnect
-		# TODO: add a timeout
-		await self._connect_result
+		try:
+			await asyncio.wait_for(self._connect_result, self.connect_timeout)
+		except asyncio.TimeoutError:
+			result = "timeout"
+		else:
 			result = self._connect_result.result()
+		
 		logger.debug(f"&{roomname}._connect_result: {result!r}")
 		
 		# deal with result