From 6141475b6aa7f2111a5d7e58ff1f044ae2419e3b Mon Sep 17 00:00:00 2001
From: Adam Dyess <adam.dyess@canonical.com>
Date: Fri, 9 Aug 2024 16:52:22 -0500
Subject: [PATCH] Handle various microcluster ready endpoints

---
 .../k8s/lib/charms/k8s/v0/k8sd_api_manager.py | 27 ++++++++--
 .../k8s/tests/unit/test_k8sd_api_manager.py   | 49 ++++++++++++++++++-
 2 files changed, 72 insertions(+), 4 deletions(-)

diff --git a/charms/worker/k8s/lib/charms/k8s/v0/k8sd_api_manager.py b/charms/worker/k8s/lib/charms/k8s/v0/k8sd_api_manager.py
index 02ca839c..7fea7ccf 100644
--- a/charms/worker/k8s/lib/charms/k8s/v0/k8sd_api_manager.py
+++ b/charms/worker/k8s/lib/charms/k8s/v0/k8sd_api_manager.py
@@ -776,9 +776,30 @@ def is_cluster_ready(self):
         return status.metadata and status.metadata.status.ready
 
     def check_k8sd_ready(self):
-        """Check if k8sd is ready."""
-        endpoint = "/cluster/1.0/ready"
-        self._send_request(endpoint, "GET", EmptyResponse)
+        """Check if k8sd is ready using various microcluster endpoints.
+
+        Raises:
+            K8sdConnectionError: If the response is Not Found on all endpoints.
+        """
+        ready_endpoints = ["/core/1.0/ready", "/cluster/1.0/ready"]
+        for endpoint in ready_endpoints:
+            try:
+                self._send_request(endpoint, "GET", EmptyResponse)
+                break
+            except InvalidResponseError as ex:
+                if ex.code == 404:
+                    logger.warning(
+                        "Encountered 404 while checking if micro-cluster is ready @ %s: %s",
+                        endpoint,
+                        ex,
+                    )
+                    # Try the next endpoint if the current one is not found
+                    continue
+                raise
+        else:
+            raise K8sdConnectionError(
+                "Exhausted all endpoints while checking if micro-cluster is ready"
+            )
 
     def bootstrap_k8s_snap(self, request: CreateClusterRequest) -> None:
         """Bootstrap the k8s cluster.
diff --git a/charms/worker/k8s/tests/unit/test_k8sd_api_manager.py b/charms/worker/k8s/tests/unit/test_k8sd_api_manager.py
index e58d49cb..b57d4c73 100644
--- a/charms/worker/k8s/tests/unit/test_k8sd_api_manager.py
+++ b/charms/worker/k8s/tests/unit/test_k8sd_api_manager.py
@@ -7,7 +7,7 @@
 import socket
 import unittest
 from socket import AF_UNIX, SOCK_STREAM
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, call, patch
 
 from lib.charms.k8s.v0.k8sd_api_manager import (
     AuthTokenResponse,
@@ -143,6 +143,53 @@ def setUp(self):
         self.mock_factory = MagicMock()
         self.api_manager = K8sdAPIManager(factory=self.mock_factory)
 
+    @patch("lib.charms.k8s.v0.k8sd_api_manager.K8sdAPIManager._send_request")
+    def test_check_k8sd_in_error(self, mock_send_request):
+        """Test bootstrap."""
+        not_found = InvalidResponseError(code=404, msg="Not Found")
+        in_error = InvalidResponseError(code=504, msg="In Error")
+        mock_send_request.side_effect = [not_found, in_error]
+
+        with self.assertRaises(InvalidResponseError) as ie:
+            self.api_manager.check_k8sd_ready()
+        mock_send_request.assert_has_calls(
+            [
+                call("/core/1.0/ready", "GET", EmptyResponse),
+                call("/cluster/1.0/ready", "GET", EmptyResponse),
+            ]
+        )
+        assert ie.exception.code == 504
+
+    @patch("lib.charms.k8s.v0.k8sd_api_manager.K8sdAPIManager._send_request")
+    def test_check_k8sd_not_found(self, mock_send_request):
+        """Test bootstrap."""
+        not_found = InvalidResponseError(code=404, msg="Not Found")
+        mock_send_request.side_effect = [not_found, not_found]
+
+        with self.assertRaises(K8sdConnectionError):
+            self.api_manager.check_k8sd_ready()
+        mock_send_request.assert_has_calls(
+            [
+                call("/core/1.0/ready", "GET", EmptyResponse),
+                call("/cluster/1.0/ready", "GET", EmptyResponse),
+            ]
+        )
+
+    @patch("lib.charms.k8s.v0.k8sd_api_manager.K8sdAPIManager._send_request")
+    def test_check_k8sd_ready(self, mock_send_request):
+        """Test bootstrap."""
+        not_found = InvalidResponseError(code=404, msg="Not Found")
+        success = EmptyResponse(status_code=200, type="test", error_code=0)
+        mock_send_request.side_effect = [not_found, success]
+
+        self.api_manager.check_k8sd_ready()
+        mock_send_request.assert_has_calls(
+            [
+                call("/core/1.0/ready", "GET", EmptyResponse),
+                call("/cluster/1.0/ready", "GET", EmptyResponse),
+            ]
+        )
+
     @patch("lib.charms.k8s.v0.k8sd_api_manager.K8sdAPIManager._send_request")
     def test_bootstrap_k8s_snap(self, mock_send_request):
         """Test bootstrap."""