Fix

Signed-off-by: Tyler Gu <[email protected]>

acto/kubernetes_engine/base.py

@@ -1,7 +1,7 @@
 import subprocess
 import time
 from abc import ABC, abstractmethod
-from typing import Callable, Dict, List
+from typing import Callable, Dict, Generator, Iterator, List
 
 import kubernetes
 
@@ -86,3 +86,11 @@ def get_node_list(self, name: str):
             # no nodes can be found, returning an empty array
             return []
         return p.stdout.strip().split('\n')
+
+    @abstractmethod
+    def stream_control_plane_stats(self, id: str) -> Iterator[Dict]:
+        '''Get stats of a node
+        Args:
+            - id: ID of the node, either name or ID
+        '''
+        raise NotImplementedError

acto/kubernetes_engine/kind.py

@@ -2,10 +2,11 @@
 import os
 import subprocess
 import time
-from typing import Dict, List
+from typing import Dict, Generator, Iterator, List
 
 import kubernetes
 import yaml
+import docker
 
 from acto.common import kubernetes_client, print_event
 from acto.constant import CONST
@@ -67,6 +68,11 @@ def get_context_name(self, cluster_name: str) -> str:
         KIND always adds `kind` before the cluster name
         '''
         return f'kind-{cluster_name}'
+
+    def get_control_plane_node_name(self, cluster_name: str) -> str:
+        '''Returns the name of the control plane node
+        '''
+        return f'{cluster_name}-control-plane'
 
     def create_cluster(self, name: str, kubeconfig: str):
         '''Use subprocess to create kind cluster
@@ -166,3 +172,21 @@ def get_node_list(self, name: str):
             raise RuntimeError
 
         return res
+
+    def stream_control_plane_stats(self, cluster_name: str) -> Generator[Dict]:
+        '''Get stats of a node
+        Args:
+            - cluster_name: name of the cluster
+        '''
+        return self.get_node_stats(self.get_control_plane_node_name(cluster_name))
+
+    def get_node_stats(self, id: str) -> Iterator[Dict]:
+        '''Get stats of a node
+        Args:
+            - id: ID of the node, either name or ID
+        '''
+        docker_client = docker.from_env()
+        container = docker_client.containers.get(id)
+        stats = container.stats(decode=True, stream=True)
+        for stat in stats:
+            yield stat