Replace AssertionError with ReconcilerError (#173)

charms/worker/k8s/src/charm.py

@@ -191,7 +191,6 @@ def _k8s_info(self, event: ops.EventBase):
     @status.on_error(
         ops.WaitingStatus("Installing COS requirements"),
         subprocess.CalledProcessError,
-        AssertionError,
     )
     def _apply_cos_requirements(self):
         """Apply COS requirements for integration.
@@ -315,7 +314,7 @@ def _check_k8sd_ready(self):
 
     @on_error(
         ops.WaitingStatus("Waiting to bootstrap k8s snap"),
-        AssertionError,
+        ReconcilerError,
         InvalidResponseError,
         K8sdConnectionError,
     )
@@ -387,7 +386,7 @@ def _get_valid_annotations(self) -> Optional[dict]:
             dict: The parsed annotations if valid, otherwise None.
 
         Raises:
-            AssertionError: If any annotation is invalid.
+            ReconcilerError: If any annotation is invalid.
         """
         raw_annotations = self.config.get("annotations")
         if not raw_annotations:
@@ -398,9 +397,10 @@ def _get_valid_annotations(self) -> Optional[dict]:
         annotations = {}
         try:
             for key, value in [pair.split("=", 1) for pair in raw_annotations.split()]:
-                assert key and value, "Invalid Annotation"  # nosec
+                if not key or not value:
+                    raise ReconcilerError("Invalid Annotation")
                 annotations[key] = value
-        except AssertionError:
+        except ReconcilerError:
             log.exception("Invalid annotations: %s", raw_annotations)
             status.add(ops.BlockedStatus("Invalid Annotations"))
             raise
@@ -456,14 +456,18 @@ def _configure_datastore(self, config: Union[BootstrapConfig, UpdateClusterConfi
                 ", ".join(SUPPORTED_DATASTORES),
             )
             status.add(ops.BlockedStatus(f"Invalid datastore: {datastore}"))
-        assert datastore in SUPPORTED_DATASTORES  # nosec
+        if datastore not in SUPPORTED_DATASTORES:
+            raise ReconcilerError(f"Invalid datastore: {datastore}")
 
         if datastore == "etcd":
             log.info("Using etcd as external datastore")
             etcd_relation = self.model.get_relation("etcd")
 
-            assert etcd_relation, "Missing etcd relation"  # nosec
-            assert self.etcd.is_ready, "etcd is not ready"  # nosec
+            if not etcd_relation:
+                raise ReconcilerError("Missing etcd relation")
+
+            if not self.etcd.is_ready:
+                raise ReconcilerError("etcd is not ready")
 
             etcd_config = self.etcd.get_client_credentials()
             if isinstance(config, BootstrapConfig):
@@ -580,7 +584,7 @@ def _enable_functionalities(self):
 
     @on_error(
         WaitingStatus("Ensure that the cluster configuration is up-to-date"),
-        AssertionError,
+        ReconcilerError,
         InvalidResponseError,
         K8sdConnectionError,
     )
@@ -618,7 +622,7 @@ def _get_scrape_jobs(self):
                 return self.cos.get_metrics_endpoints(
                     self.get_node_name(), token, self.is_control_plane
                 )
-        except AssertionError:
+        except ReconcilerError:
             log.exception("Failed to get COS token.")
         return []
 
@@ -690,7 +694,7 @@ def _get_proxy_env(self) -> Dict[str, str]:
 
     @on_error(
         WaitingStatus("Waiting for Cluster token"),
-        AssertionError,
+        ReconcilerError,
         InvalidResponseError,
         K8sdConnectionError,
     )

charms/worker/k8s/src/token_distributor.py

@@ -11,6 +11,7 @@
 
 import charms.contextual_status as status
 import ops
+from charms.contextual_status import ReconcilerError
 from charms.k8s.v0.k8sd_api_manager import (
     ErrorCodes,
     InvalidResponseError,
@@ -209,6 +210,9 @@ def cluster_name(self, relation: ops.Relation, local: bool) -> str:
 
         Returns:
             the recovered cluster name from existing relations
+
+        Raises:
+            ReconcilerError: If fails to find 1 relation-name:cluster-name.
         """
         cluster_name: Optional[str] = ""
         if not local:
@@ -218,7 +222,8 @@ def cluster_name(self, relation: ops.Relation, local: bool) -> str:
                 if value := relation.data[unit].get("cluster-name"):
                     values |= {value}
             if values:
-                assert len(values) == 1, f"Failed to find 1 {relation.name}:cluster-name"  # nosec
+                if len(values) != 1:
+                    raise ReconcilerError(f"Failed to find 1 {relation.name}:cluster-name")
                 (cluster_name,) = values
         elif not (cluster_name := relation.data[self.charm.unit].get("joined")):
             # joined_cluster_name
@@ -235,6 +240,12 @@ def recover_token(self, relation: ops.Relation) -> Generator[str, None, None]:
 
         Yields:
             str: extracted token content
+
+        Raises:
+            ReconcilerError:
+                - If fails to find 1 relation-name:secret-id.
+                - If relation-name:secret-key is not valid.
+                - If relation-name:token is not valid.
         """
         self.request(relation)
 
@@ -246,14 +257,17 @@ def recover_token(self, relation: ops.Relation) -> Generator[str, None, None]:
             if (secret_id := relation.data[unit].get(secret_key))
         }
 
-        assert len(secret_ids) == 1, f"Failed to find 1 {relation.name}:{secret_key}"  # nosec
+        if len(secret_ids) != 1:
+            raise ReconcilerError(f"Failed to find 1 {relation.name}:{secret_key}")
         (secret_id,) = secret_ids
-        assert secret_id, f"{relation.name}:{secret_key} is not valid"  # nosec
+        if not secret_id:
+            raise ReconcilerError(f"{relation.name}:{secret_key} is not valid")
         secret = self.charm.model.get_secret(id=secret_id)
 
         # Get the content from the secret
         content = secret.get_content(refresh=True)
-        assert content["token"], f"{relation.name}:token not valid"  # nosec
+        if not content.get("token"):
+            raise ReconcilerError(f"{relation.name}:token not valid")
         yield content["token"]
 
         # signal that the relation is joined, the token is used
@@ -343,7 +357,7 @@ def update_node(self, relation: ops.Relation, unit: ops.Unit, state: str):
         """
         relation.data[self.charm.app][unit.name] = state
 
-    def allocate_tokens(
+    def allocate_tokens(  # noqa: C901
         self,
         relation: ops.Relation,
         token_strategy: TokenStrategy,
@@ -356,16 +370,23 @@ def allocate_tokens(
             token_strategy (TokenStrategy): The strategy of token creation.
             token_type (ClusterTokenType): The type of cluster token.
                 Defaults to ClusterTokenType.NONE.
+
+        Raises:
+            ReconcilerError:
+                - If token_strategy is valid.
+                - If remote application doesn't exist on relation.
         """
         units = relation.units
         if self.charm.app == relation.app:
             # include self in peer relations
             units |= {self.charm.unit}
-        assert relation.app, f"Remote application doesn't exist on {relation.name}"  # nosec
+        if not relation.app:
+            raise ReconcilerError(f"Remote application doesn't exist on {relation.name}")
 
         # Select the appropriate token creation strategy
         tokenizer = self.token_strategies.get(token_strategy)
-        assert tokenizer, f"Invalid token_strategy: {token_strategy}"  # nosec
+        if not tokenizer:
+            raise ReconcilerError(f"Invalid token_strategy: {token_strategy}")
 
         log.info("Allocating %s %s tokens", token_type.name.title(), token_strategy.name.title())
         status.add(