diff --git a/acto/kubernetes_engine/base.py b/acto/kubernetes_engine/base.py index e3312b2a91..f9f48c1057 100644 --- 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 index 3de948a7c2..f26ed60edf 100644 --- a/acto/kubernetes_engine/kind.py +++ b/acto/kubernetes_engine/kind.py @@ -2,11 +2,12 @@ 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 @@ -14,8 +15,11 @@ 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']