Add support for multiple endpoints in Connection

Author: David Garcia

juju/client/connection.py

@@ -230,7 +230,7 @@ async def connect(
         If uuid is None, the connection will be to the controller. Otherwise it
         will be to the model.
 
-        :param str endpoint: The hostname:port of the controller to connect to.
+        :param str endpoint: The hostname:port of the controller to connect to (or list of strings).
         :param str uuid: The model UUID to connect to (None for a
             controller-only connection).
         :param str username: The username for controller-local users (or None
@@ -256,6 +256,8 @@ async def connect(
         self = cls()
         if endpoint is None:
             raise ValueError('no endpoint provided')
+        if not isinstance(endpoint, str) and not isinstance(endpoint, list):
+            raise TypeError("Endpoint should be either str or list")
         self.uuid = uuid
         if bakery_client is None:
             bakery_client = httpbakery.Client()
@@ -279,6 +281,7 @@ async def connect(
         self.addr = None
         self.ws = None
         self.endpoint = None
+        self.endpoints = None
         self.cacert = None
         self.info = None
 
@@ -297,7 +300,11 @@ async def connect(
         if max_frame_size is None:
             max_frame_size = self.MAX_FRAME_SIZE
         self.max_frame_size = max_frame_size
-        await self._connect_with_redirect([(endpoint, cacert)])
+        await self._connect_with_redirect(
+            [(endpoint, cacert)]
+            if isinstance(endpoint, str)
+            else [(e, cacert) for e in endpoint]
+        )
         return self
 
     @property
@@ -570,7 +577,11 @@ async def reconnect(self):
             return
         async with monitor.reconnecting:
             await self.close()
-            await self._connect_with_login([(self.endpoint, self.cacert)])
+            await self._connect_with_login(
+                [(self.endpoint, self.cacert)]
+                if not self.endpoints else
+                self.endpoints
+            )
 
     async def _connect(self, endpoints):
         if len(endpoints) == 0:
@@ -660,6 +671,8 @@ async def _connect_with_login(self, endpoints):
         finally:
             if not success:
                 await self.close()
+            else:
+                self._pinger_task.start()
 
     async def _connect_with_redirect(self, endpoints):
         try:

juju/client/connector.py

@@ -67,6 +67,13 @@ async def connect(self, **kwargs):
             for macaroon in kwargs.pop('macaroons'):
                 jar.set_cookie(go_to_py_cookie(macaroon))
         self._connection = await Connection.connect(**kwargs)
+        controller = self.jujudata.controllers()[
+            self.jujudata.current_controller()
+        ]
+        self._connection.endpoints = [
+            (e, controller["ca-cert"])
+            for e in controller["api-endpoints"]
+        ]
 
     async def disconnect(self):
         """Shut down the watcher task and close websockets.
@@ -86,13 +93,11 @@ async def connect_controller(self, controller_name=None, specified_facades=None)
             raise JujuConnectionError('No current controller')
 
         controller = self.jujudata.controllers()[controller_name]
-        # TODO change Connection so we can pass all the endpoints
-        # instead of just the first.
-        endpoint = controller['api-endpoints'][0]
+        endpoints = controller['api-endpoints']
         accounts = self.jujudata.accounts().get(controller_name, {})
 
         await self.connect(
-            endpoint=endpoint,
+            endpoint=endpoints,
             uuid=None,
             username=accounts.get('user'),
             password=accounts.get('password'),
@@ -119,9 +124,7 @@ async def connect_model(self, model_name=None):
         if controller is None:
             raise JujuConnectionError('Controller {} not found'.format(
                 controller_name))
-        # TODO change Connection so we can pass all the endpoints
-        # instead of just the first one.
-        endpoint = controller['api-endpoints'][0]
+        endpoints = controller['api-endpoints']
         account = self.jujudata.accounts().get(controller_name, {})
         models = self.jujudata.models().get(controller_name, {}).get('models',
                                                                      {})
@@ -135,7 +138,7 @@ async def connect_model(self, model_name=None):
         # and also remove the need for base.CleanModel to
         # subclass JujuData.
         await self.connect(
-            endpoint=endpoint,
+            endpoint=endpoints,
             uuid=models[model_name]['uuid'],
             username=account.get('user'),
             password=account.get('password'),