From 41351d786171a0747d8ee0f1fb1b933d8d99f946 Mon Sep 17 00:00:00 2001
From: Caner Derici <caner.derici@canonical.com>
Date: Fri, 13 Oct 2023 14:42:21 -0600
Subject: [PATCH 1/2] Improve jujudata to find controller name by endpoint

Fixes #771
---
 juju/client/connector.py | 13 +++++++++----
 juju/client/jujudata.py  | 20 +++++++++++++++++++-
 juju/errors.py           |  2 ++
 3 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/juju/client/connector.py b/juju/client/connector.py
index 7499de66..4747b3d7 100644
--- a/juju/client/connector.py
+++ b/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',
                                                                      {})
diff --git a/juju/client/jujudata.py b/juju/client/jujudata.py
index 79e924c6..ce9f5ae9 100644
--- a/juju/client/jujudata.py
+++ b/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')
 
diff --git a/juju/errors.py b/juju/errors.py
index a34dfbd0..08831a6e 100644
--- a/juju/errors.py
+++ b/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):

From 7162c6d399b89544ec4139190d7ce77ae1b533b7 Mon Sep 17 00:00:00 2001
From: Caner Derici <caner.derici@canonical.com>
Date: Fri, 13 Oct 2023 14:49:30 -0600
Subject: [PATCH 2/2] Fix linter

---
 juju/client/jujudata.py | 1 +
 juju/errors.py          | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/juju/client/jujudata.py b/juju/client/jujudata.py
index ce9f5ae9..67976e60 100644
--- a/juju/client/jujudata.py
+++ b/juju/client/jujudata.py
@@ -15,6 +15,7 @@
 
 API_ENDPOINTS_KEY = 'api-endpoints'
 
+
 class NoModelException(Exception):
     pass
 
diff --git a/juju/errors.py b/juju/errors.py
index 08831a6e..2232420c 100644
--- a/juju/errors.py
+++ b/juju/errors.py
@@ -89,9 +89,11 @@ class JujuUnitError(JujuError):
 class JujuBackupError(JujuError):
     pass
 
+
 class PylibjujuProgrammingError(Exception):
     pass
 
+
 class JujuNotValid(JujuError):
     def __init__(self, entity_type, entity_name):
         self.entity_type = entity_type