Feat: support posthooks for KubernetesEngine

Signed-off-by: Tyler Gu <[email protected]>
MarkintoshZ · Oct 1, 2023 · c6ff460 · c6ff460
1 parent 64cab37
commit c6ff460
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/acto/kubernetes_engine/base.py b/acto/kubernetes_engine/base.py
@@ -1,12 +1,22 @@
 import subprocess
 import time
 from abc import ABC, abstractmethod
+from typing import Callable, List
+
+import kubernetes
 
 from acto.constant import CONST
 from acto.utils import get_thread_logger
 
+KubernetesEnginePostHookType = Callable[[kubernetes.client.ApiClient], None]
+
 
 class KubernetesEngine(ABC):
+
+    @abstractmethod
+    def __init__(self, acto_namespace: int,
+                 posthooks: List[KubernetesEnginePostHookType] = None) -> None: ...
+
     @abstractmethod
     def configure_cluster(self, num_nodes: int, version: str):
         pass
@@ -29,7 +39,7 @@ def delete_cluster(self, name: str, kubeconfig: str, ):
 
     def restart_cluster(self, name: str, kubeconfig: str):
         logger = get_thread_logger(with_prefix=False)
-        
+
         retry_count = 3
 
         while (retry_count > 0):

diff --git a/acto/kubernetes_engine/kind.py b/acto/kubernetes_engine/kind.py
@@ -2,20 +2,24 @@
 import os
 import subprocess
 import time
+from typing import List
 
 import kubernetes
 import yaml
-from acto.common import print_event
 
+from acto.common import kubernetes_client, print_event
 from acto.constant import CONST
 
 from . import base
 
 
 class Kind(base.KubernetesEngine):
 
-    def __init__(self, acto_namespace: int):
+    def __init__(
+            self, acto_namespace: int, posthooks: List[base.KubernetesEnginePostHookType] = None):
         self.config_path = os.path.join(CONST.CLUSTER_CONFIG_FOLDER, f'KIND-{acto_namespace}.yaml')
+        if posthooks is not None:
+            self.posthooks = posthooks
 
     def configure_cluster(self, num_nodes: int, version: str):
         '''Create config file for kind'''
@@ -93,12 +97,17 @@ def create_cluster(self, name: str, kubeconfig: str):
         try:
             kubernetes.config.load_kube_config(config_file=kubeconfig,
                                                context=self.get_context_name(name))
+            apiclient = kubernetes_client(kubeconfig, self.get_context_name(name))
         except Exception as e:
             logging.debug("Incorrect kube config file:")
             with open(kubeconfig) as f:
                 logging.debug(f.read())
             raise e
 
+        if self.posthooks:
+            for posthook in self.posthooks:
+                posthook(apiclient=apiclient)
+
     def load_images(self, images_archive_path: str, name: str):
         logging.info('Loading preload images')
         cmd = ['kind', 'load', 'image-archive']