From d7d7b9f4307830bdfe0117a7523ffa14f22da0c4 Mon Sep 17 00:00:00 2001
From: Jordan Barrett <jordan.barrett@canonical.com>
Date: Wed, 8 Nov 2023 16:12:52 +0200
Subject: [PATCH] use urllib to parse apiaddress - test ipv4 and ipv6

---
 src/charm.py        |  6 +++++-
 tests/test_charm.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/src/charm.py b/src/charm.py
index 4b2c580..302014b 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -5,6 +5,7 @@
 import controlsocket
 import logging
 import secrets
+import urllib.parse
 import yaml
 from charms.prometheus_k8s.v0.prometheus_scrape import MetricsEndpointProvider
 from ops.charm import CharmBase
@@ -129,7 +130,10 @@ def api_port(self) -> str:
         if not isinstance(api_addresses, List):
             raise AgentConfException("agent.conf key 'apiaddresses' is not a list")
 
-        return api_addresses[0].split(':')[-1]
+        parsed_url = urllib.parse.urlsplit('//' + api_addresses[0])
+        if not parsed_url.port:
+            raise AgentConfException("api address doesn't include port")
+        return parsed_url.port
 
     def ca_cert(self) -> str:
         """Return the controller's CA certificate."""
diff --git a/tests/test_charm.py b/tests/test_charm.py
index a90febe..3d51315 100644
--- a/tests/test_charm.py
+++ b/tests/test_charm.py
@@ -24,6 +24,18 @@
 cacert: fake
 '''
 
+agent_conf_ipv4 = '''
+apiaddresses:
+- "127.0.0.1:17070"
+cacert: fake
+'''
+
+agent_conf_ipv6 = '''
+apiaddresses:
+- "[::1]:17070"
+cacert: fake
+'''
+
 
 class TestCharm(unittest.TestCase):
     def test_relation_joined(self):
@@ -132,6 +144,22 @@ def test_apiaddresses_missing_status(self, *_):
             "can't read controller API port from agent.conf: agent.conf key 'apiaddresses' missing"
         ))
 
+    @patch("builtins.open", new_callable=mock_open, read_data=agent_conf_ipv4)
+    def test_apiaddresses_ipv4(self, _):
+        harness = Harness(JujuControllerCharm)
+        self.addCleanup(harness.cleanup)
+        harness.begin()
+
+        self.assertEqual(harness.charm.api_port(), 17070)
+
+    @patch("builtins.open", new_callable=mock_open, read_data=agent_conf_ipv6)
+    def test_apiaddresses_ipv6(self, _):
+        harness = Harness(JujuControllerCharm)
+        self.addCleanup(harness.cleanup)
+        harness.begin()
+
+        self.assertEqual(harness.charm.api_port(), 17070)
+
 
 class MockBinding:
     def __init__(self, address):