From d227579d56f827a96ac8563f7820e7960f3a92a1 Mon Sep 17 00:00:00 2001 From: mathieui Date: Sun, 24 Jan 2021 21:33:14 +0100 Subject: xmlstream: set disconnected future on event --- slixmpp/xmlstream/xmlstream.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) 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) -- cgit v1.2.3