From d98bc6570b7af1a3e437992e05382a6bc992b75e Mon Sep 17 00:00:00 2001
From: Tyler Gu <jiaweig3@illinois.edu>
Date: Thu, 13 Jun 2024 22:08:57 -0500
Subject: [PATCH] Support num nodes for kubernetes

Signed-off-by: Tyler Gu <jiaweig3@illinois.edu>
---
 acto/kubernetes_engine/kind.py        |  2 +-
 chactos/fault_injection_config.py     | 17 +++++++++++++++++
 chactos/fault_injections.py           |  4 +++-
 chactos/percona-mongodb-operator.json |  5 ++++-
 4 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/acto/kubernetes_engine/kind.py b/acto/kubernetes_engine/kind.py
index 1b348a4f8b..ee5c1d00cf 100644
--- a/acto/kubernetes_engine/kind.py
+++ b/acto/kubernetes_engine/kind.py
@@ -39,7 +39,7 @@ def __init__(
         extra_mounts.append(
             {"hostPath": "profile/data", "containerPath": "/tmp/profile"}
         )
-        for _ in range(num_nodes - 1):
+        for _ in range(num_nodes):
             config_dict["nodes"].append(
                 {
                     "role": "worker",
diff --git a/chactos/fault_injection_config.py b/chactos/fault_injection_config.py
index f5f56da2b5..33d29084f1 100644
--- a/chactos/fault_injection_config.py
+++ b/chactos/fault_injection_config.py
@@ -1,8 +1,24 @@
+from typing import Optional
+
 import pydantic
 
 from acto.lib.operator_config import DeployConfig
 
 
+class KubernetesConfig(pydantic.BaseModel, extra="forbid"):
+    """Kubernetes Config"""
+
+    num_nodes: int = pydantic.Field(
+        description="Number of workers in the Kubernetes cluster", default=4
+    )
+    version: str = pydantic.Field(
+        default="v1.28.0", description="Kubernetes version"
+    )
+    feature_gates: Optional[dict[str, bool]] = pydantic.Field(
+        description="Path to the feature gates file", default=None
+    )
+
+
 class FaultInjectionConfig(pydantic.BaseModel, extra="forbid"):
     """Fault Injection Config"""
 
@@ -11,3 +27,4 @@ class FaultInjectionConfig(pydantic.BaseModel, extra="forbid"):
     operator_selector: dict
     application_pod_prefix: str
     input_dir: str
+    kubernetes: KubernetesConfig
diff --git a/chactos/fault_injections.py b/chactos/fault_injections.py
index 41bbfa42f9..38880b432f 100644
--- a/chactos/fault_injections.py
+++ b/chactos/fault_injections.py
@@ -43,7 +43,9 @@ def __init__(self, operator_config: FaultInjectionConfig, worker_id: int):
 
     def run(self):
         """Run the experiment."""
-        k8s_cluster_engine: KubernetesEngine = Kind(0)
+        k8s_cluster_engine: KubernetesEngine = Kind(
+            0, num_nodes=self._operator_config.num_nodes
+        )
         cluster_name = f"acto-cluster-{self._worker_id}"
         kubecontext = k8s_cluster_engine.get_context_name(cluster_name)
         kubeconfig = os.path.join(os.path.expanduser("~"), ".kube", kubecontext)
diff --git a/chactos/percona-mongodb-operator.json b/chactos/percona-mongodb-operator.json
index b75bcececa..36d10ef4bd 100644
--- a/chactos/percona-mongodb-operator.json
+++ b/chactos/percona-mongodb-operator.json
@@ -16,5 +16,8 @@
         "labelSelectors": {"name": "percona-server-mongodb-operator"}
     },
     "application_pod_prefix": "test-cluster-rs0-",
-    "input_dir": "data/percona-server-mongodb-operator/v1.16.0/inputs"
+    "input_dir": "data/percona-server-mongodb-operator/v1.16.0/inputs",
+    "kubernetes": {
+        "num_nodes": 7
+    }
 }
\ No newline at end of file