Add Driver.disconnect function

src/mattermostdriver/driver.py

@@ -148,6 +148,10 @@ async def my_event_handler(message):
 		loop.run_until_complete(self.websocket.connect(event_handler))
 		return loop
 
+	def disconnect(self):
+		"""Disconnects the driver from the server, stopping the websocket event loop."""
+		self.websocket.disconnect()
+
 	def login(self):
 		"""
 		Logs the user in.
@@ -418,4 +422,3 @@ def integration_actions(self):
 		:return: Instance of :class:`~endpoints.integration_actions.IntegrationActions`
 		"""
 		return IntegrationActions(self.client)
-

src/mattermostdriver/websocket.py

@@ -14,6 +14,7 @@ def __init__(self, options, token):
 		if options['debug']:
 			log.setLevel(logging.DEBUG)
 		self._token = token
+		self._alive = False
 
 	async def connect(self, event_handler):
 		"""
@@ -56,7 +57,8 @@ async def _start_loop(self, websocket, event_handler):
 		forcing us to reconnect.
 		"""
 		log.debug('Starting websocket loop')
-		while True:
+		self._alive = True
+		while self._alive:
 			try:
 				await asyncio.wait_for(
 					self._wait_for_message(websocket, event_handler),
@@ -67,6 +69,11 @@ async def _start_loop(self, websocket, event_handler):
 				log.debug("Sending heartbeat...")
 				continue
 
+	def disconnect(self):
+		"""Sets `self._alive` to False so the loop in `self._start_loop` will finish."""
+		log.info("Disconnecting websocket")
+		self._alive = False
+
 	async def _authenticate_websocket(self, websocket, event_handler):
 		"""
 		Sends a authentication challenge over a websocket.