Improve jujudata to find controller name by endpoint

Fixes #771

Author: cderici

juju/client/connector.py

@@ -7,9 +7,9 @@
 import macaroonbakery.httpbakery as httpbakery
 from juju.client.connection import Connection
 from juju.client.gocookies import GoCookieJar, go_to_py_cookie
-from juju.client.jujudata import FileJujuData
+from juju.client.jujudata import FileJujuData, API_ENDPOINTS_KEY
 from juju.client.proxy.factory import proxy_from_config
-from juju.errors import JujuConnectionError, JujuError
+from juju.errors import JujuConnectionError, JujuError, PylibjujuProgrammingError
 from juju.client import client
 from juju.version import SUPPORTED_MAJOR_VERSION, TARGET_JUJU_VERSION
 
@@ -83,6 +83,11 @@ async def connect(self, **kwargs):
                 await self._connection.close()
             self._connection = await Connection.connect(**kwargs)
 
+        if not self.controller_name:
+            if 'endpoint' not in kwargs:
+                raise PylibjujuProgrammingError("Please report this error to the maintainers.")
+            self.controller_name = self.jujudata.controller_name_by_endpoint(kwargs['endpoint'])
+
         # Check if we support the target controller
         juju_server_version = self._connection.info['server-version']
         if not juju_server_version.startswith(TARGET_JUJU_VERSION):
@@ -112,7 +117,7 @@ async def connect_controller(self, controller_name=None, specified_facades=None)
             raise JujuConnectionError('No current controller')
 
         controller = self.jujudata.controllers()[controller_name]
-        endpoints = controller['api-endpoints']
+        endpoints = controller[API_ENDPOINTS_KEY]
         accounts = self.jujudata.accounts().get(controller_name, {})
 
         proxy = proxy_from_config(controller.get('proxy-config', None))
@@ -146,7 +151,7 @@ async def connect_model(self, _model_name=None, **kwargs):
         if controller is None:
             raise JujuConnectionError('Controller {} not found'.format(
                 controller_name))
-        endpoints = controller['api-endpoints']
+        endpoints = controller[API_ENDPOINTS_KEY]
         account = self.jujudata.accounts().get(controller_name, {})
         models = self.jujudata.models().get(controller_name, {}).get('models',
                                                                      {})

juju/client/jujudata.py

@@ -10,9 +10,10 @@
 import yaml
 from juju import tag
 from juju.client.gocookies import GoCookieJar
-from juju.errors import JujuError
+from juju.errors import JujuError, PylibjujuProgrammingError
 from juju.utils import juju_config_dir
 
+API_ENDPOINTS_KEY = 'api-endpoints'
 
 class NoModelException(Exception):
     pass
@@ -126,6 +127,23 @@ def load_credential(self, cloud, name=None):
         except (KeyError, FileNotFoundError):
             return None, None
 
+    def controller_name_by_endpoint(self, endpoint):
+        """Finds the controller that has the given endpoints, returns the name.
+
+        :param str endpoint: The endpoint of the controller we're looking for
+        """
+        for controller_name, controller in self.controllers().items():
+            if isinstance(endpoint, str):
+                if endpoint in controller[API_ENDPOINTS_KEY]:
+                    return controller_name
+            elif isinstance(endpoint, list):
+                for e in endpoint:
+                    if e in controller[API_ENDPOINTS_KEY]:
+                        return controller_name
+            else:
+                raise PylibjujuProgrammingError()
+        raise JujuError(f'Unable to find controller with endpoint {endpoint}')
+
     def controllers(self):
         return self._load_yaml('controllers.yaml', 'controllers')
 

juju/errors.py

@@ -89,6 +89,8 @@ class JujuUnitError(JujuError):
 class JujuBackupError(JujuError):
     pass
 
+class PylibjujuProgrammingError(Exception):
+    pass
 
 class JujuNotValid(JujuError):
     def __init__(self, entity_type, entity_name):