44
55from ocp_resources .pod import Pod
66from ocp_utilities .exceptions import CommandExecFailed
7- from timeout_sampler import TimeoutExpiredError , TimeoutSampler , retry
7+ from timeout_sampler import retry
88
99from libs .vm .vm import BaseVirtualMachine
1010
1515LOGGER = logging .getLogger (__name__ )
1616
1717
18- class BaseClient (ABC ):
18+ class BaseTcpClient (ABC ):
1919 """Base abstract class for network traffic generator client."""
2020
2121 def __init__ (self , server_ip : str , server_port : int ):
@@ -24,7 +24,7 @@ def __init__(self, server_ip: str, server_port: int):
2424 self ._cmd = f"{ _IPERF_BIN } --client { self ._server_ip } --time 0 --port { self .server_port } --connect-timeout 300"
2525
2626 @abstractmethod
27- def __enter__ (self ) -> "BaseClient " :
27+ def __enter__ (self ) -> "BaseTcpClient " :
2828 pass
2929
3030 @abstractmethod
@@ -36,7 +36,7 @@ def is_running(self) -> bool:
3636 pass
3737
3838
39- class Server :
39+ class TcpServer :
4040 """
4141 Represents a server running on a virtual machine for testing network performance.
4242 Implemented with iperf3
@@ -55,11 +55,13 @@ def __init__(
5555 self ._port = port
5656 self ._cmd = f"{ _IPERF_BIN } --server --port { self ._port } --one-off"
5757
58- def __enter__ (self ) -> "Server " :
58+ def __enter__ (self ) -> "TcpServer " :
5959 self ._vm .console (
6060 commands = [f"{ self ._cmd } &" ],
6161 timeout = _DEFAULT_CMD_TIMEOUT_SEC ,
6262 )
63+ self ._ensure_is_running ()
64+
6365 return self
6466
6567 def __exit__ (self , exc_type : BaseException , exc_value : BaseException , traceback : object ) -> None :
@@ -72,10 +74,13 @@ def vm(self) -> BaseVirtualMachine:
7274 def is_running (self ) -> bool :
7375 return _is_process_running (vm = self ._vm , cmd = self ._cmd )
7476
77+ @retry (wait_timeout = 30 , sleep = 2 , exceptions_dict = {})
78+ def _ensure_is_running (self ) -> bool :
79+ return self .is_running ()
80+
7581
76- class Client (BaseClient ):
77- """
78- Represents a client that connects to a server to test network performance.
82+ class VMTcpClient (BaseTcpClient ):
83+ """Represents a TCP client that connects to a server to test network performance.
7984 Implemented with iperf3
8085
8186 Args:
@@ -93,11 +98,13 @@ def __init__(
9398 super ().__init__ (server_ip = server_ip , server_port = server_port )
9499 self ._vm = vm
95100
96- def __enter__ (self ) -> "Client " :
101+ def __enter__ (self ) -> "VMTcpClient " :
97102 self ._vm .console (
98103 commands = [f"{ self ._cmd } &" ],
99104 timeout = _DEFAULT_CMD_TIMEOUT_SEC ,
100105 )
106+ self ._ensure_is_running ()
107+
101108 return self
102109
103110 def __exit__ (self , exc_type : BaseException , exc_value : BaseException , traceback : object ) -> None :
@@ -110,6 +117,10 @@ def vm(self) -> BaseVirtualMachine:
110117 def is_running (self ) -> bool :
111118 return _is_process_running (vm = self ._vm , cmd = self ._cmd )
112119
120+ @retry (wait_timeout = 30 , sleep = 2 , exceptions_dict = {})
121+ def _ensure_is_running (self ) -> bool :
122+ return self .is_running ()
123+
113124
114125def _stop_process (vm : BaseVirtualMachine , cmd : str ) -> None :
115126 try :
@@ -118,25 +129,18 @@ def _stop_process(vm: BaseVirtualMachine, cmd: str) -> None:
118129 LOGGER .warning (str (e ))
119130
120131
121- def _is_process_running ( # type: ignore[return]
122- vm : BaseVirtualMachine , cmd : str
123- ) -> bool :
132+ def _is_process_running (vm : BaseVirtualMachine , cmd : str ) -> bool :
124133 try :
125- for sample in TimeoutSampler (
126- wait_timeout = 60 ,
127- sleep = 5 ,
128- func = vm .console ,
134+ vm .console (
129135 commands = [f"pgrep -fx '{ cmd } '" ],
130136 timeout = _DEFAULT_CMD_TIMEOUT_SEC ,
131- ):
132- if sample :
133- return True
134- except TimeoutExpiredError as e :
135- LOGGER .warning (f"Process is not running on VM { vm .name } . Error: { str (e .last_exp )} " )
137+ )
138+ return True
139+ except CommandExecFailed :
136140 return False
137141
138142
139- class PodClient ( BaseClient ):
143+ class PodTcpClient ( BaseTcpClient ):
140144 """Represents a TCP client that connects to a server to test network performance.
141145
142146 Expects pod to have iperf3 container.
@@ -155,7 +159,7 @@ def __init__(self, pod: Pod, server_ip: str, server_port: int, bind_interface: s
155159 self ._container = _IPERF_BIN
156160 self ._cmd += f" --bind { bind_interface } " if bind_interface else ""
157161
158- def __enter__ (self ) -> "PodClient " :
162+ def __enter__ (self ) -> "PodTcpClient " :
159163 # run the command in the background using nohup to ensure it keeps running after the exec session ends
160164 self ._pod .execute (
161165 command = ["sh" , "-c" , f"nohup { self ._cmd } >/tmp/{ _IPERF_BIN } .log 2>&1 &" ], container = self ._container
@@ -178,5 +182,5 @@ def _ensure_is_running(self) -> bool:
178182 return self .is_running ()
179183
180184
181- def is_tcp_connection (server : Server , client : Client | PodClient ) -> bool :
185+ def is_tcp_connection (server : TcpServer , client : BaseTcpClient ) -> bool :
182186 return server .is_running () and client .is_running ()
0 commit comments