From d5b349e03106cfede3eddd4fc1444d13f21b7cd8 Mon Sep 17 00:00:00 2001
From: Dan Elkouby <streetwalkermc@gmail.com>
Date: Sun, 17 Jun 2018 20:29:41 +0300
Subject: [PATCH] Implement a mechanism to notify of connection failures (#852)

---
 telethon/client/telegrambaseclient.py |  8 ++++++++
 telethon/client/updates.py            | 13 +++++++++++++
 telethon/network/mtprotosender.py     | 26 +++++++++++++++++++++++---
 3 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/telethon/client/telegrambaseclient.py b/telethon/client/telegrambaseclient.py
index 1bc002fe..866a756a 100644
--- a/telethon/client/telegrambaseclient.py
+++ b/telethon/client/telegrambaseclient.py
@@ -218,6 +218,14 @@ class TelegramBaseClient(abc.ABC):
     def loop(self):
         return self._loop
 
+    @property
+    def disconnected(self):
+        """
+        Future that resolves when the connection to Telegram
+        ends, either by user action or in the background.
+        """
+        return self._sender.disconnected
+
     # endregion
 
     # region Connecting
diff --git a/telethon/client/updates.py b/telethon/client/updates.py
index 72d63928..4685444e 100644
--- a/telethon/client/updates.py
+++ b/telethon/client/updates.py
@@ -14,6 +14,19 @@ class UpdateMethods(UserMethods):
 
     # region Public methods
 
+    def run_until_disconnected(self):
+        """
+        Runs the event loop until `disconnect` is called or if an error
+        while connecting/sending/receiving occurs in the background. In
+        the latter case, said error will ``raise`` so you have a chance
+        to ``except`` it on your own code.
+
+        This method shouldn't be called from ``async def`` as the loop
+        will be running already. Use ``await client.disconnected`` in
+        this situation instead.
+        """
+        self.loop.run_until_complete(self.disconnected)
+
     def on(self, event):
         """
         Decorator helper method around add_event_handler().
diff --git a/telethon/network/mtprotosender.py b/telethon/network/mtprotosender.py
index 438bb8bd..f27ee5ef 100644
--- a/telethon/network/mtprotosender.py
+++ b/telethon/network/mtprotosender.py
@@ -63,6 +63,7 @@ class MTProtoSender:
         # pending futures should be cancelled.
         self._user_connected = False
         self._reconnecting = False
+        self._disconnected = None
 
         # We need to join the loops upon disconnection
         self._send_loop_handle = None
@@ -157,6 +158,10 @@ class MTProtoSender:
             self._recv_loop_handle.cancel()
 
         __log__.info('Disconnection from {} complete!'.format(self._ip))
+        if error:
+            self._disconnected.set_exception(error)
+        else:
+            self._disconnected.set_result(None)
 
     def send(self, request, ordered=False):
         """
@@ -199,6 +204,17 @@ class MTProtoSender:
             self._send_queue.put_nowait(message)
             return message.future
 
+    @property
+    def disconnected(self):
+        """
+        Future that resolves when the connection to Telegram
+        ends, either by user action or in the background.
+        """
+        if self._disconnected is not None:
+            return self._disconnected
+        else:
+            raise ConnectionError('Sender was never connected')
+
     # Private methods
 
     async def _connect(self):
@@ -235,9 +251,10 @@ class MTProtoSender:
                 else:
                     break
             else:
-                await self._disconnect()
-                raise ConnectionError('auth_key generation failed {} times'
-                                      .format(self._retries))
+                e = ConnectionError('auth_key generation failed {} times'
+                                    .format(self._retries))
+                await self._disconnect(error=e)
+                raise e
 
         __log__.debug('Starting send loop')
         self._send_loop_handle = self._loop.create_task(self._send_loop())
@@ -245,6 +262,9 @@ class MTProtoSender:
         __log__.debug('Starting receive loop')
         self._recv_loop_handle = self._loop.create_task(self._recv_loop())
 
+        # First connection or manual reconnection after a failure
+        if self._disconnected is None or self._disconnected.done():
+            self._disconnected = asyncio.Future()
         __log__.info('Connection to {} complete!'.format(self._ip))
 
     async def _reconnect(self):