Implement XWaitFor for Features (#458)

Add `x-wait-for` command to allow the snap to wait until features are available and ready in the cluster.

src/k8s/cmd/k8s/k8s.go

@@ -94,6 +94,7 @@ func NewRootCmd(env cmdutil.ExecutionEnvironment) *cobra.Command {
 		newHelmCmd(env),
 		xPrintShimPidsCmd,
 		newXSnapdConfigCmd(env),
+		newXWaitForCmd(env),
 	)
 
 	cmd.DisableAutoGenTag = true

src/k8s/cmd/k8s/k8s_x_wait_for.go

@@ -0,0 +1,49 @@
+package k8s
+
+import (
+	cmdutil "github.com/canonical/k8s/cmd/util"
+	"github.com/canonical/k8s/pkg/k8sd/features"
+	"github.com/canonical/k8s/pkg/utils/control"
+	"github.com/spf13/cobra"
+)
+
+func newXWaitForCmd(env cmdutil.ExecutionEnvironment) *cobra.Command {
+	waitForDNSCmd := &cobra.Command{
+		Use:   "dns",
+		Short: "Wait for DNS to be ready",
+		Run: func(cmd *cobra.Command, args []string) {
+			err := control.WaitUntilReady(cmd.Context(), func() (bool, error) {
+				return features.StatusChecks.CheckDNS(cmd.Context(), env.Snap)
+			})
+			if err != nil {
+				cmd.PrintErrf("Error: failed to wait for DNS to be ready: %v\n", err)
+				env.Exit(1)
+			}
+		},
+	}
+
+	waitForNetworkCmd := &cobra.Command{
+		Use:   "network",
+		Short: "Wait for Network to be ready",
+		Run: func(cmd *cobra.Command, args []string) {
+			err := control.WaitUntilReady(cmd.Context(), func() (bool, error) {
+				return features.StatusChecks.CheckNetwork(cmd.Context(), env.Snap)
+			})
+			if err != nil {
+				cmd.PrintErrf("Error: failed to wait for DNS to be ready: %v\n", err)
+				env.Exit(1)
+			}
+		},
+	}
+
+	cmd := &cobra.Command{
+		Use:    "x-wait-for",
+		Short:  "Wait for the cluster's feature to be in a ready state",
+		Hidden: true,
+	}
+
+	cmd.AddCommand(waitForDNSCmd)
+	cmd.AddCommand(waitForNetworkCmd)
+
+	return cmd
+}

src/k8s/pkg/client/kubernetes/pods.go

@@ -0,0 +1,34 @@
+package kubernetes
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+// IsPodReady checks if a pod is ready.
+func (c *Client) IsPodReady(ctx context.Context, name, namespace string, listOptions metav1.ListOptions) (bool, error) {
+	pods, err := c.CoreV1().Pods(namespace).List(ctx, listOptions)
+	if err != nil {
+		return false, fmt.Errorf("failed to list pods: %w", err)
+	}
+
+	for _, pod := range pods.Items {
+		if strings.Contains(pod.Name, name) {
+			if pod.Status.Phase != corev1.PodRunning {
+				return false, nil
+			}
+
+			for _, condition := range pod.Status.Conditions {
+				if condition.Type == corev1.PodReady && condition.Status == corev1.ConditionTrue {
+					return true, nil
+				}
+			}
+		}
+	}
+
+	return false, nil
+}

src/k8s/pkg/k8sd/features/cilium/status.go

@@ -0,0 +1,34 @@
+package cilium
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/canonical/k8s/pkg/snap"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func CheckNetwork(ctx context.Context, snap snap.Snap) (bool, error) {
+	client, err := snap.KubernetesClient("kube-system")
+	if err != nil {
+		return false, fmt.Errorf("failed to create kubernetes client: %w", err)
+	}
+
+	ciliumPods := map[string]string{
+		"cilium-operator": "io.cilium/app=operator",
+		"cilium":          "k8s-app=cilium",
+	}
+
+	for ciliumPod, selector := range ciliumPods {
+		isReady, err := client.IsPodReady(ctx, ciliumPod, "kube-system", metav1.ListOptions{LabelSelector: selector})
+		if err != nil {
+			return false, fmt.Errorf("failed to check if pod %q is ready: %w", ciliumPod, err)
+		}
+		if !isReady {
+			return false, nil
+		}
+	}
+
+	return true, nil
+}

src/k8s/pkg/k8sd/features/coredns/status.go

@@ -0,0 +1,25 @@
+package coredns
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/canonical/k8s/pkg/snap"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+// CheckDNS checks the CoreDNS deployment in the cluster.
+func CheckDNS(ctx context.Context, snap snap.Snap) (bool, error) {
+	client, err := snap.KubernetesClient("kube-system")
+	if err != nil {
+		return false, fmt.Errorf("failed to create kubernetes client: %w", err)
+	}
+
+	isReady, err := client.IsPodReady(ctx, "coredns", "kube-system", metav1.ListOptions{LabelSelector: "app.kubernetes.io/name=coredns"})
+	if err != nil {
+		return false, fmt.Errorf("failed to wait for CoreDNS pod to be ready: %w", err)
+	}
+
+	return isReady, nil
+}

src/k8s/pkg/k8sd/features/implementation_default.go

@@ -21,3 +21,9 @@ var Implementation Interface = &implementation{
 	applyMetricsServer: metrics_server.ApplyMetricsServer,
 	applyLocalStorage:  localpv.ApplyLocalStorage,
 }
+
+// StatusChecks implements the Canonical Kubernetes built-in feature status checks.
+var StatusChecks StatusInterface = &statusChecks{
+	checkNetwork: cilium.CheckNetwork,
+	checkDNS:     coredns.CheckDNS,
+}

src/k8s/pkg/k8sd/features/status.go

@@ -0,0 +1,25 @@
+package features
+
+import (
+	"context"
+
+	"github.com/canonical/k8s/pkg/snap"
+)
+
+type StatusInterface interface {
+	CheckDNS(context.Context, snap.Snap) (bool, error)
+	CheckNetwork(context.Context, snap.Snap) (bool, error)
+}
+
+type statusChecks struct {
+	checkDNS     func(context.Context, snap.Snap) (bool, error)
+	checkNetwork func(context.Context, snap.Snap) (bool, error)
+}
+
+func (s *statusChecks) CheckDNS(ctx context.Context, snap snap.Snap) (bool, error) {
+	return s.checkDNS(ctx, snap)
+}
+
+func (s *statusChecks) CheckNetwork(ctx context.Context, snap snap.Snap) (bool, error) {
+	return s.checkNetwork(ctx, snap)
+}

tests/integration/tests/test_util/util.py

@@ -153,69 +153,13 @@ def wait_until_k8s_ready(
 
 
 def wait_for_dns(instance: harness.Instance):
-    LOG.info("Waiting for CoreDNS pod to show up...")
-    stubbornly(retries=15, delay_s=5).on(instance).until(
-        lambda p: "coredns" in p.stdout.decode()
-    ).exec(["k8s", "kubectl", "get", "pod", "-n", "kube-system", "-o", "json"])
-    LOG.info("CoreDNS pod showed up.")
-
-    stubbornly(retries=3, delay_s=1).on(instance).exec(
-        [
-            "k8s",
-            "kubectl",
-            "wait",
-            "--for=condition=ready",
-            "pod",
-            "-n",
-            "kube-system",
-            "-l",
-            "app.kubernetes.io/name=coredns",
-            "--timeout",
-            "180s",
-        ]
-    )
+    LOG.info("Waiting for DNS to be ready")
+    instance.exec(["k8s", "x-wait-for", "dns"])
 
 
 def wait_for_network(instance: harness.Instance):
-    LOG.info("Waiting for cilium pods to show up...")
-    stubbornly(retries=15, delay_s=5).on(instance).until(
-        lambda p: "cilium" in p.stdout.decode()
-    ).exec(
-        ["k8s", "kubectl", "get", "pod", "-n", "kube-system", "-o", "json"],
-    )
-    LOG.info("Cilium pods showed up.")
-
-    stubbornly(retries=3, delay_s=1).on(instance).exec(
-        [
-            "k8s",
-            "kubectl",
-            "wait",
-            "--for=condition=ready",
-            "pod",
-            "-n",
-            "kube-system",
-            "-l",
-            "io.cilium/app=operator",
-            "--timeout",
-            "180s",
-        ],
-    )
-
-    stubbornly(retries=3, delay_s=1).on(instance).exec(
-        [
-            "k8s",
-            "kubectl",
-            "wait",
-            "--for=condition=ready",
-            "pod",
-            "-n",
-            "kube-system",
-            "-l",
-            "k8s-app=cilium",
-            "--timeout",
-            "180s",
-        ]
-    )
+    LOG.info("Waiting for network to be ready")
+    instance.exec(["k8s", "x-wait-for", "network"])
 
 
 def hostname(instance: harness.Instance) -> str: