Close the internally opened connection on the controller

This connection is opened to get the model_uuid, and creates its own
pinger/receiver tasks etc that need to be gracefully handled before
the main model is disconnected at teardown.

Author: cderici

juju/client/connector.py

@@ -73,6 +73,12 @@ async def connect(self, **kwargs):
             assert self._connection
             self._log_connection = await Connection.connect(**kwargs)
         else:
+            # TODO (cderici): we need to investigate how to reuse/share
+            # connections between Model & Controller objects
+            # At least either avoid overwriting self._connection with
+            # different types of connections (model, controller), or
+            # have an indication of which type of entity this is
+            # connected to.
             if self._connection:
                 await self._connection.close()
             self._connection = await Connection.connect(**kwargs)
@@ -84,11 +90,11 @@ async def connect(self, **kwargs):
         if not self._connection.info['server-version'].startswith(SUPPORTED_MAJOR_VERSION):
             raise JujuConnectionError("juju server-version %s not supported" % juju_server_version)
 
-    async def disconnect(self):
+    async def disconnect(self, entity):
         """Shut down the watcher task and close websockets.
         """
         if self._connection:
-            log.debug('Closing model connection')
+            log.debug(f'Closing {entity} connection')
             await self._connection.close()
             self._connection = None
         if self._log_connection:

juju/controller.py

@@ -193,7 +193,7 @@ async def disconnect(self):
         """Shut down the watcher task and close websockets.
 
         """
-        await self._connector.disconnect()
+        await self._connector.disconnect(entity='controller')
 
     async def add_credential(self, name=None, credential=None, cloud=None,
                              owner=None, force=False):

juju/model.py

@@ -772,9 +772,13 @@ async def watch_received_waiter():
             raise self._watcher_task.exception()
 
         if self.info is None:
-            contr = await self.get_controller()
-            self._info = await contr.get_model_info(model_name, model_uuid)
-            log.debug('Got ModelInfo: %s', vars(self.info))
+            # TODO (cderici): See if this can be optimized away, or at least
+            # be done lazily (i.e. not everytime after_connect, but whenever
+            # self.info is needed -- which here can be bypassed if model_uuid
+            # is known)
+            async with ConnectedController(self.connection()) as contr:
+                self._info = await contr.get_model_info(model_name, model_uuid)
+                log.debug('Got ModelInfo: %s', vars(self.info))
 
         self.uuid = self.info.uuid
 
@@ -795,8 +799,7 @@ async def disconnect(self):
             self._watch_stopping.clear()
 
         if self.is_connected():
-            log.debug('Closing model connection')
-            await self._connector.disconnect()
+            await self._connector.disconnect(entity='Model')
             self._info = None
 
     async def add_local_charm_dir(self, charm_dir, series):