From a0317c1d2647146118ed32d60f9ec74d9dfbbe17 Mon Sep 17 00:00:00 2001
From: Xudong Sun <xudongs3@illinois.edu>
Date: Tue, 12 Sep 2023 22:30:31 -0500
Subject: [PATCH] Refactor zk e2e tests'

Signed-off-by: Xudong Sun <xudongs3@illinois.edu>
---
 e2e/src/common.rs        |  8 ++++++++
 e2e/src/zookeeper_e2e.rs | 39 ++++++++++++---------------------------
 2 files changed, 20 insertions(+), 27 deletions(-)

diff --git a/e2e/src/common.rs b/e2e/src/common.rs
index d2ff86de5..4640c5f84 100644
--- a/e2e/src/common.rs
+++ b/e2e/src/common.rs
@@ -12,7 +12,9 @@ use kube::{
     discovery::{ApiCapabilities, ApiResource, Discovery, Scope},
     Client, CustomResource,
 };
+use std::process::Command;
 use thiserror::Error;
+
 #[derive(Debug, Error)]
 pub enum Error {
     #[error("Failed to get kube client: {0}")]
@@ -120,3 +122,9 @@ pub async fn get_output_and_err(mut attached: AttachedProcess) -> (String, Strin
     attached.join().await.unwrap();
     (out, err)
 }
+
+pub fn run_command(program: &str, args: Vec<&str>, err_msg: &str) {
+    let cmd = Command::new(program).args(args).output().expect(err_msg);
+    println!("cmd output: {}", String::from_utf8_lossy(&cmd.stdout));
+    println!("cmd error: {}", String::from_utf8_lossy(&cmd.stderr));
+}
diff --git a/e2e/src/zookeeper_e2e.rs b/e2e/src/zookeeper_e2e.rs
index ea5bb7fff..274b54333 100644
--- a/e2e/src/zookeeper_e2e.rs
+++ b/e2e/src/zookeeper_e2e.rs
@@ -1,7 +1,7 @@
 #![allow(unused_imports)]
 #![allow(unused_variables)]
-use k8s_openapi::api::core::v1::{ConfigMap, Pod};
-use k8s_openapi::api::{apps::v1::StatefulSet, core::v1::Service};
+use k8s_openapi::api::apps::v1::StatefulSet;
+use k8s_openapi::api::core::v1::{ConfigMap, Pod, Service};
 use k8s_openapi::apiextensions_apiserver::pkg::apis::apiextensions::v1::CustomResourceDefinition;
 use kube::{
     api::{
@@ -16,7 +16,6 @@ use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use serde_json::json;
 use std::path::PathBuf;
-use std::process::Command;
 use std::thread;
 use std::time::{Duration, Instant};
 use tokio::time::sleep;
@@ -159,24 +158,17 @@ pub async fn scaling_test(client: Client, zk_name: String) -> Result<(), Error>
     let timeout = Duration::from_secs(360);
     let start = Instant::now();
     let sts_api: Api<StatefulSet> = Api::default_namespaced(client.clone());
-    let scale_output = Command::new("kubectl")
-        .args([
+    run_command(
+        "kubectl",
+        vec![
             "patch",
             "zk",
             "zookeeper",
             "--type=json",
             "-p",
             "[{\"op\": \"replace\", \"path\": \"/spec/replicas\", \"value\": 2}]",
-        ])
-        .output()
-        .expect("failed to scale zk");
-    println!(
-        "cmd output: {}",
-        String::from_utf8_lossy(&scale_output.stdout)
-    );
-    println!(
-        "cmd error: {}",
-        String::from_utf8_lossy(&scale_output.stderr)
+        ],
+        "failed to scale zk",
     );
 
     loop {
@@ -228,24 +220,17 @@ pub async fn scaling_test(client: Client, zk_name: String) -> Result<(), Error>
         };
     }
 
-    Command::new("kubectl")
-        .args([
+    run_command(
+        "kubectl",
+        vec![
             "patch",
             "zk",
             "zookeeper",
             "--type=json",
             "-p",
             "[{\"op\": \"replace\", \"path\": \"/spec/replicas\", \"value\": 3}]",
-        ])
-        .output()
-        .expect("failed to scale zk");
-    println!(
-        "cmd output: {}",
-        String::from_utf8_lossy(&scale_output.stdout)
-    );
-    println!(
-        "cmd error: {}",
-        String::from_utf8_lossy(&scale_output.stderr)
+        ],
+        "failed to scale zk",
     );
 
     loop {