Test helper to create a temporary k3d cluster for kube-rs

kube_test_k3d.rs

use std::{convert::TryFrom, process::Command};

use kube::{config::Kubeconfig, Client, Config};

#[derive(Debug, Default)]
pub struct TestEnv {
    // Unique name.
    name: String,
    // Kubeconfig of the temporary cluster.
    kubeconfig: Kubeconfig,
}

impl TestEnv {
    /// Builder for configuring the test environemnt.
    pub fn builder() -> TestEnvBuilder {
        Default::default()
    }

    /// Create the default minimal test environment.
    pub fn new() -> Self {
        Self::builder().build()
    }

    fn delete(&mut self) {
        println!("Deleting k3d cluster {}...", &self.name);
        let status = Command::new("k3d")
            .args(&["cluster", "delete", &self.name])
            .status()
            .expect("k3d cluster delete failed");
        assert!(
            status.success(),
            "k3d cluster delete failed. cluster {} may still exist",
            self.name
        );
    }

    /// Create a new `Client` configured for the temporary server.
    pub async fn client(&self) -> Client {
        assert_eq!(
            self.kubeconfig.clusters.len(),
            1,
            "kubeconfig only contains the temporary cluster"
        );
        assert_eq!(
            self.kubeconfig
                .clusters
                .get(0)
                .unwrap()
                .name
                .as_str()
                .strip_prefix("k3d-")
                .unwrap(),
            self.name,
            "kubeconfig only contains the temporary cluster"
        );

        let config = Config::from_custom_kubeconfig(self.kubeconfig.clone(), &Default::default())
            .await
            .expect("valid kubeconfig");
        Client::try_from(config).expect("client")
    }
}

impl Drop for TestEnv {
    fn drop(&mut self) {
        self.delete();
    }
}


#[derive(Debug)]
pub struct TestEnvBuilder {
    /// The number of servers. Default: 1
    servers: usize,
    /// The number of agents. Default: 0
    agents: usize,
    /// Inject the Host IP as `host.k3d.internal` into the containers and CoreDNS.
    /// Default: false
    host_ip_injection: bool,
    /// Create an image volume for importing images.
    /// Default: false
    create_image_volume: bool,
    /// Create a `LoadBalancer` in front of the server nodes.
    /// Default: false
    create_load_balancer: bool,
    /// Set `--verbose` flag. Default: false
    verbose: bool,
}

impl Default for TestEnvBuilder {
    fn default() -> Self {
        Self {
            servers: 1,
            agents: 0,
            host_ip_injection: false,
            create_image_volume: false,
            create_load_balancer: false,
            verbose: false,
        }
    }
}

impl TestEnvBuilder {
    pub fn servers(&mut self, servers: usize) -> &mut Self {
        self.servers = servers;
        self
    }

    pub fn agents(&mut self, agents: usize) -> &mut Self {
        self.agents = agents;
        self
    }

    pub fn inject_host_ip(&mut self) -> &mut Self {
        self.host_ip_injection = true;
        self
    }

    pub fn with_image_volume(&mut self) -> &mut Self {
        self.create_image_volume = true;
        self
    }

    pub fn with_load_balancer(&mut self) -> &mut Self {
        self.create_load_balancer = true;
        self
    }

    pub fn verbose(&mut self) -> &mut Self {
        self.verbose = true;
        self
    }

    pub fn build(&self) -> TestEnv {
        let name = xid::new().to_string();
        let servers = format!("--servers={}", self.servers);
        let agents = format!("--agents={}", self.agents);
        let mut args = vec![
            "cluster",
            "create",
            &name,
            "--wait",
            // Don't change `~/.kube/config`
            "--kubeconfig-update-default=false",
            "--kubeconfig-switch-context=false",
            // Disable to avoid having to create the default service account in each test.
            "--k3s-server-arg",
            "--kube-apiserver-arg=disable-admission-plugins=ServiceAccount",
            // Disable components and features
            "--k3s-server-arg",
            "--disable=servicelb",
            "--k3s-server-arg",
            "--disable=traefik",
            "--k3s-server-arg",
            "--disable=metrics-server",
            "--k3s-server-arg",
            "--disable-cloud-controller",
            "--no-rollback",
            &servers,
            &agents,
        ];
        if self.verbose {
            args.push("--verbose");
        }
        if !self.host_ip_injection {
            args.push("--no-hostip");
        }
        if !self.create_image_volume {
            args.push("--no-image-volume");
        }
        if !self.create_load_balancer {
            args.push("--no-lb");
        }

        let status = Command::new("k3d")
            .args(&args)
            .status()
            .expect("k3d cluster create");
        assert!(status.success(), "failed to create k3d cluster");

        // Output the cluster's kubeconfig to stdout and store it.
        let stdout = Command::new("k3d")
            .args(&["kubeconfig", "get", &name])
            .output()
            .expect("k3d kubeconfig get failed")
            .stdout;
        let stdout = std::str::from_utf8(&stdout).expect("valid string");

        TestEnv {
            name,
            kubeconfig: serde_yaml::from_str(stdout).expect("valid kubeconfig"),
        }
    }
}


#[tokio::test]
async fn test_integration() {
    use k8s_openapi::api::core::v1::Pod;
    use kube::Api;

    let test_env = TestEnv::new();
    let client = test_env.client().await;

    let pods: Api<Pod> = Api::default_namespaced(client);
    let _pod = pods.get("example").await.unwrap();
}