Avoid gathering all tasks when reconnect-ing

Author: cderici

juju/client/connection.py

@@ -218,7 +218,8 @@ def status(self):
         if connection.is_debug_log_connection:
             stopped = connection._debug_log_task.cancelled()
         else:
-            stopped = connection._receiver_task.cancelled()
+            stopped = connection._receiver_task is not None and \
+                      connection._receiver_task.cancelled()
 
         if stopped or not connection._ws.open:
             return self.ERROR
@@ -451,17 +452,17 @@ async def close(self, to_reconnect=False):
         if self._ws and not self._ws.closed:
             await self._ws.close()
 
-        try:
-            log.debug('Gathering all tasks for connection close')
+        if not to_reconnect:
+            try:
+                log.debug('Gathering all tasks for connection close')
 
-            # Avoid gathering the current task
-            tasks_need_to_be_gathered = [task for task in jasyncio.all_tasks()
-                                         if task != jasyncio.current_task()]
-            await jasyncio.gather(*tasks_need_to_be_gathered)
-        except jasyncio.CancelledError:
-            pass
-        except websockets.exceptions.ConnectionClosed:
-            pass
+                # Avoid gathering the current task
+                tasks_need_to_be_gathered = [task for task in jasyncio.all_tasks() if task != jasyncio.current_task()]
+                await jasyncio.gather(*tasks_need_to_be_gathered)
+            except jasyncio.CancelledError:
+                pass
+            except websockets.exceptions.ConnectionClosed:
+                pass
 
         self._pinger_task = None
         self._receiver_task = None

tests/integration/test_connection.py

@@ -31,7 +31,7 @@ async def test_monitor(event_loop):
         assert conn.monitor.status == 'connected'
         await conn.close()
 
-        assert conn.monitor.status == 'disconnected'
+        assert conn.monitor.status == 'disconnecting'
 
 
 @base.bootstrapped