11import contextlib
22import logging
33from abc import ABC , abstractmethod
4- from typing import Final , Generator
4+ from collections .abc import Generator
5+ from typing import Final
56
67from ocp_resources .pod import Pod
78from ocp_utilities .exceptions import CommandExecFailed
@@ -32,7 +33,7 @@ def server_ip(self) -> str:
3233 return self ._server_ip
3334
3435 @abstractmethod
35- def __enter__ (self ) -> " BaseTcpClient" :
36+ def __enter__ (self ) -> BaseTcpClient :
3637 pass
3738
3839 @abstractmethod
@@ -66,7 +67,7 @@ def __init__(
6667 self ._cmd = f"{ _IPERF_BIN } --server --port { self ._port } --one-off"
6768 self ._cmd += f" --bind { bind_ip } " if bind_ip else ""
6869
69- def __enter__ (self ) -> " TcpServer" :
70+ def __enter__ (self ) -> TcpServer :
7071 self ._vm .console (
7172 commands = [f"{ self ._cmd } &" ],
7273 timeout = _DEFAULT_CMD_TIMEOUT_SEC ,
@@ -113,7 +114,7 @@ def __init__(
113114 self ._vm = vm
114115 self ._cmd += f" --set-mss { maximum_segment_size } " if maximum_segment_size else ""
115116
116- def __enter__ (self ) -> " VMTcpClient" :
117+ def __enter__ (self ) -> VMTcpClient :
117118 self ._vm .console (
118119 commands = [f"{ self ._cmd } &" ],
119120 timeout = _DEFAULT_CMD_TIMEOUT_SEC ,
@@ -174,7 +175,7 @@ def __init__(self, pod: Pod, server_ip: str, server_port: int, bind_interface: s
174175 self ._container = _IPERF_BIN
175176 self ._cmd += f" --bind { bind_interface } " if bind_interface else ""
176177
177- def __enter__ (self ) -> " PodTcpClient" :
178+ def __enter__ (self ) -> PodTcpClient :
178179 # run the command in the background using nohup to ensure it keeps running after the exec session ends
179180 self ._pod .execute (
180181 command = ["sh" , "-c" , f"nohup { self ._cmd } >/tmp/{ _IPERF_BIN } .log 2>&1 &" ], container = self ._container
@@ -204,7 +205,7 @@ def active_tcp_connections(
204205 client_vm : BaseVirtualMachine ,
205206 server_vm : BaseVirtualMachine ,
206207 iface_name : str ,
207- ) -> Generator [list [tuple [VMTcpClient , TcpServer ]], None , None ]:
208+ ) -> Generator [list [tuple [VMTcpClient , TcpServer ]]]:
208209 """Start iperf3 client-server connections for all IPs on the server's interface.
209210 The helper assumed the ip addresses are up.
210211
@@ -242,7 +243,7 @@ def client_server_active_connection(
242243 port : int = IPERF_SERVER_PORT ,
243244 maximum_segment_size : int = 0 ,
244245 ip_family : int = 4 ,
245- ) -> Generator [tuple [VMTcpClient , TcpServer ], None , None ]:
246+ ) -> Generator [tuple [VMTcpClient , TcpServer ]]:
246247 """Start iperf3 client-server connection with continuous TCP traffic flow.
247248
248249 Automatically starts an iperf3 server and client, with traffic flowing continuously
@@ -265,11 +266,13 @@ def client_server_active_connection(
265266 Traffic runs with infinite duration until context exits.
266267 """
267268 server_ip = str (lookup_iface_status_ip (vm = server_vm , iface_name = spec_logical_network , ip_family = ip_family ))
268- with TcpServer (vm = server_vm , port = port , bind_ip = server_ip ) as server :
269- with VMTcpClient (
269+ with (
270+ TcpServer (vm = server_vm , port = port , bind_ip = server_ip ) as server ,
271+ VMTcpClient (
270272 vm = client_vm ,
271273 server_ip = server_ip ,
272274 server_port = port ,
273275 maximum_segment_size = maximum_segment_size ,
274- ) as client :
275- yield client , server
276+ ) as client ,
277+ ):
278+ yield client , server
0 commit comments