Safely configure datastore request in the same method (#82)

* Safely configure datastore request in the same method

* Update charm.py

If etcd isn't ready via AssertionError -- just wait til it's ready

* With unit testing, found we were passing the wrong object into _configure_datastore

charms/worker/k8s/src/charm.py

@@ -24,7 +24,7 @@
 from functools import cached_property
 from pathlib import Path
 from time import sleep
-from typing import Dict, Optional
+from typing import Dict, Optional, Union
 from urllib.parse import urlparse
 
 import charms.contextual_status as status
@@ -264,12 +264,12 @@ def _configure_cos_integration(self):
         if relation := self.model.get_relation("cos-tokens"):
             self.collector.request(relation)
 
-    def _configure_datastore(self, config: BootstrapConfig):
+    def _configure_datastore(self, config: Union[BootstrapConfig, UpdateClusterConfigRequest]):
         """Configure the datastore for the Kubernetes cluster.
 
         Args:
-            config (BootstrapConfig): The bootstrap configuration object for
-                the Kubernetes cluster that is being configured. This object
+            config (BootstrapConfig|UpdateClusterConfigRequst):
+                The configuration object for the Kubernetes cluster. This object
                 will be modified in-place to include etcd's configuration details.
         """
         datastore = self.config.get("datastore")
@@ -290,12 +290,24 @@ def _configure_datastore(self, config: BootstrapConfig):
             assert etcd_relation, "Missing etcd relation"  # nosec
             assert self.etcd.is_ready, "etcd is not ready"  # nosec
 
-            config.datastore_type = "external"
             etcd_config = self.etcd.get_client_credentials()
-            config.datastore_ca_cert = etcd_config.get("client_ca", "")
-            config.datastore_client_cert = etcd_config.get("client_cert", "")
-            config.datastore_client_key = etcd_config.get("client_key", "")
-            config.datastore_servers = self.etcd.get_connection_string().split(",")
+            if isinstance(config, BootstrapConfig):
+                config.datastore_type = "external"
+                config.datastore_ca_cert = etcd_config.get("client_ca", "")
+                config.datastore_client_cert = etcd_config.get("client_cert", "")
+                config.datastore_client_key = etcd_config.get("client_key", "")
+                config.datastore_servers = self.etcd.get_connection_string().split(",")
+                log.info("etcd servers: %s", config.datastore_servers)
+            elif isinstance(config, UpdateClusterConfigRequest):
+                config.datastore = UserFacingDatastoreConfig(
+                    type="external",
+                    servers=self.etcd.get_connection_string().split(","),
+                    ca_crt=etcd_config.get("client_ca", ""),
+                    client_crt=etcd_config.get("client_cert", ""),
+                    client_key=etcd_config.get("client_key", ""),
+                )
+                log.info("etcd servers: %s", config.datastore.servers)
+
         elif datastore == "dqlite":
             log.info("Using dqlite as datastore")
 
@@ -390,6 +402,7 @@ def _enable_functionalities(self):
 
     @on_error(
         WaitingStatus("Ensure that the cluster configuration is up-to-date"),
+        AssertionError,
         InvalidResponseError,
         K8sdConnectionError,
     )
@@ -404,18 +417,7 @@ def _ensure_cluster_config(self):
 
         update_request = UpdateClusterConfigRequest()
 
-        # TODO: Ensure other configs here as well.
-
-        if self.config.get("datastore") == "etcd":
-            etcd_config = self.etcd.get_client_credentials()
-            update_request.datastore = UserFacingDatastoreConfig(
-                type="external",
-                servers=self.etcd.get_connection_string().split(","),
-                ca_crt=etcd_config.get("client_ca", ""),
-                client_crt=etcd_config.get("client_cert", ""),
-                client_key=etcd_config.get("client_key", ""),
-            )
-
+        self._configure_datastore(update_request)
         self.api_manager.update_cluster_config(update_request)
 
     def _get_scrape_jobs(self):

charms/worker/k8s/tests/unit/test_base.py

@@ -8,12 +8,14 @@
 
 
 import contextlib
+from pathlib import Path
 from unittest import mock
 
 import ops
 import ops.testing
 import pytest
 from charm import K8sCharm
+from charms.k8s.v0.k8sd_api_manager import BootstrapConfig, UpdateClusterConfigRequest
 
 
 @pytest.fixture(params=["worker", "control-plane"])
@@ -23,7 +25,10 @@ def harness(request):
     Args:
         request: pytest request object
     """
-    harness = ops.testing.Harness(K8sCharm)
+    meta = Path(__file__).parent / "../../charmcraft.yaml"
+    if request.param == "worker":
+        meta = Path(__file__).parent / "../../../charmcraft.yaml"
+    harness = ops.testing.Harness(K8sCharm, meta=meta.read_text())
     harness.begin()
     harness.charm.is_worker = request.param == "worker"
     yield harness
@@ -106,3 +111,83 @@ def test_set_leader(harness):
     assert harness.charm.reconciler.stored.reconciled
     called = {name: h for name, h in handlers.items() if h.called}
     assert len(called) == len(handlers)
+
+
+def test_configure_datastore_bootstrap_config_dqlite(harness):
+    """Test configuring the datastore=dqlite on bootstrap.
+
+    Args:
+        harness: the harness under test
+    """
+    if harness.charm.is_worker:
+        pytest.skip("Not applicable on workers")
+
+    bs_config = BootstrapConfig()
+    harness.charm._configure_datastore(bs_config)
+    assert bs_config.datastore_ca_cert is None
+    assert bs_config.datastore_client_cert is None
+    assert bs_config.datastore_client_key is None
+    assert bs_config.datastore_servers is None
+    assert bs_config.datastore_type is None
+
+
+def test_configure_datastore_bootstrap_config_etcd(harness):
+    """Test configuring the datastore=etcd on bootstrap.
+
+    Args:
+        harness: the harness under test
+    """
+    if harness.charm.is_worker:
+        pytest.skip("Not applicable on workers")
+
+    bs_config = BootstrapConfig()
+    harness.update_config({"datastore": "etcd"})
+    harness.add_relation("etcd", "etcd")
+    with mock.patch.object(harness.charm, "etcd") as mock_etcd:
+        mock_etcd.is_ready = True
+        mock_etcd.get_client_credentials.return_value = {}
+        mock_etcd.get_connection_string.return_value = "foo:1234,bar:1234"
+        harness.charm._configure_datastore(bs_config)
+    assert bs_config.datastore_ca_cert == ""
+    assert bs_config.datastore_client_cert == ""
+    assert bs_config.datastore_client_key == ""
+    assert bs_config.datastore_servers == ["foo:1234", "bar:1234"]
+    assert bs_config.datastore_type == "external"
+
+
+def test_configure_datastore_runtime_config_dqlite(harness):
+    """Test configuring the datastore=dqlite on runtime changes.
+
+    Args:
+        harness: the harness under test
+    """
+    if harness.charm.is_worker:
+        pytest.skip("Not applicable on workers")
+
+    uccr_config = UpdateClusterConfigRequest()
+    harness.charm._configure_datastore(uccr_config)
+    assert uccr_config.datastore is None
+
+
+def test_configure_datastore_runtime_config_etcd(harness):
+    """Test configuring the datastore=etcd on runtime changes.
+
+    Args:
+        harness: the harness under test
+    """
+    if harness.charm.is_worker:
+        pytest.skip("Not applicable on workers")
+
+    harness.update_config({"datastore": "etcd"})
+    harness.add_relation("etcd", "etcd")
+    with mock.patch.object(harness.charm, "etcd") as mock_etcd:
+        mock_etcd.is_ready = True
+        mock_etcd.get_client_credentials.return_value = {}
+        mock_etcd.get_connection_string.return_value = "foo:1234,bar:1234"
+        uccr_config = UpdateClusterConfigRequest()
+        harness.charm._configure_datastore(uccr_config)
+    assert uccr_config.datastore.ca_crt == ""
+    assert uccr_config.datastore.client_crt == ""
+    assert uccr_config.datastore.client_key == ""
+    assert uccr_config.datastore.servers == ["foo:1234", "bar:1234"]
+    assert uccr_config.datastore.type == "external"