diff options
-rw-r--r-- | slixmpp/xmlstream/xmlstream.py | 10 |
1 files changed, 9 insertions, 1 deletions
diff --git a/slixmpp/xmlstream/xmlstream.py b/slixmpp/xmlstream/xmlstream.py index c9247cec..5f675e3f 100644 --- a/slixmpp/xmlstream/xmlstream.py +++ b/slixmpp/xmlstream/xmlstream.py @@ -232,7 +232,7 @@ class XMLStream(asyncio.BaseProtocol): self.disconnect_reason = None #: An asyncio Future being done when the stream is disconnected. - self.disconnected = asyncio.Future() + self.disconnected: Future = Future() self.add_event_handler('disconnected', self._remove_schedules) self.add_event_handler('session_start', self._start_keepalive) @@ -259,6 +259,12 @@ class XMLStream(asyncio.BaseProtocol): """ return uuid.uuid4().hex + def _set_disconnected_future(self): + """Set the self.disconnected future on disconnect""" + if not self.disconnected.done(): + self.disconnected.set_result(True) + self.disconnected = asyncio.Future() + def connect(self, host='', port=0, use_ssl=False, force_starttls=True, disable_starttls=False): """Create a new socket and connect to the server. @@ -477,6 +483,7 @@ class XMLStream(asyncio.BaseProtocol): if self.end_session_on_disconnect: self._reset_sendq() self.event('session_end') + self._set_disconnected_future() self.event("disconnected", self.disconnect_reason or exception and exception.strerror) def cancel_connection_attempt(self): @@ -527,6 +534,7 @@ class XMLStream(asyncio.BaseProtocol): loop=self.loop, ) else: + self._set_disconnected_future() self.event("disconnected", reason) future = Future() future.set_result(None) |