Fixing typing for IPerfClient
Change-Id: I5816b3de3ed68838974c16d08508b19f7bc003e8
Reviewed-on: https://fuchsia-review.googlesource.com/c/antlion/+/865458
Fuchsia-Auto-Submit: Sam Balana <sbalana@google.com>
Commit-Queue: Auto-Submit <auto-submit@fuchsia-infra.iam.gserviceaccount.com>
Reviewed-by: Kevin Sakuma <sakuma@google.com>
diff --git a/packages/antlion/controllers/iperf_client.py b/packages/antlion/controllers/iperf_client.py
index 88b11a2..6f91d20 100644
--- a/packages/antlion/controllers/iperf_client.py
+++ b/packages/antlion/controllers/iperf_client.py
@@ -19,6 +19,7 @@
import socket
import subprocess
import threading
+from typing import Dict, List, Optional, Union
from antlion import context, utils
from antlion.controllers.adb_lib.error import AdbCommandError
@@ -47,7 +48,7 @@
Args:
configs: config parameters for the iperf server
"""
- results = []
+ results: List[IPerfClientBase] = []
for c in configs:
if type(c) is dict and "AndroidDevice" in c:
results.append(
@@ -176,12 +177,17 @@
def __init__(
self,
- ssh_config: str,
- test_interface: str = None,
- ssh_provider: SSHProvider = None,
+ ssh_config: Optional[Dict[str, Union[str, int]]] = None,
+ test_interface: Optional[str] = None,
+ ssh_provider: Optional[SSHProvider] = None,
):
self._ssh_provider = ssh_provider
+ self._ssh_session: Optional[connection.SshConnection] = None
+ self.test_interface = test_interface
+
if not self._ssh_provider:
+ if not ssh_config:
+ raise TypeError("ssh_config is required when ssh_provider is None")
self._ssh_settings = settings.from_config(ssh_config)
if not (
utils.is_valid_ipv4_address(self._ssh_settings.hostname)
@@ -192,10 +198,7 @@
)
if mdns_ip:
self._ssh_settings.hostname = mdns_ip
- self._ssh_session = None
- self.start_ssh()
-
- self.test_interface = test_interface
+ self._ssh_session = connection.SshConnection(self._ssh_settings)
def start(self, ip, iperf_args, tag, timeout=3600, iperf_binary=None):
"""Starts iperf client, and waits for completion.
@@ -223,10 +226,11 @@
full_out_path = self._get_full_file_path(tag)
try:
- self.start_ssh()
if self._ssh_provider:
iperf_process = self._ssh_provider.run(iperf_cmd, timeout_sec=timeout)
else:
+ if not self._ssh_session:
+ self._ssh_session = connection.SshConnection(self._ssh_settings)
iperf_process = self._ssh_session.run(iperf_cmd, timeout=timeout)
iperf_output = iperf_process.stdout
with open(full_out_path, "w") as out_file:
@@ -240,14 +244,6 @@
return full_out_path
- def start_ssh(self):
- """Starts an ssh session to the iperf client."""
- if self._ssh_provider:
- # SSH sessions are created by the provider.
- return
- if not self._ssh_session:
- self._ssh_session = connection.SshConnection(self._ssh_settings)
-
def close_ssh(self):
"""Closes the ssh session to the iperf client, if one exists, preventing
connection reset errors when rebooting client device.
diff --git a/packages/antlion/test_utils/abstract_devices/wmm_transceiver.py b/packages/antlion/test_utils/abstract_devices/wmm_transceiver.py
index 2506034..5fc06b4 100644
--- a/packages/antlion/test_utils/abstract_devices/wmm_transceiver.py
+++ b/packages/antlion/test_utils/abstract_devices/wmm_transceiver.py
@@ -18,7 +18,9 @@
import multiprocessing
import time
from datetime import datetime
-from uuid import uuid4
+from multiprocessing.managers import DictProxy
+from typing import Any, Optional
+from uuid import UUID, uuid4
from antlion import signals, tracelogger, utils
from antlion.controllers import iperf_client, iperf_server
@@ -396,16 +398,16 @@
def _run_traffic(
self,
- uuid,
- client,
- server_ip,
- server_port,
- active_streams,
- stream_results,
- access_category=None,
- bandwidth=None,
- stream_time=DEFAULT_STREAM_TIME,
- start_time=None,
+ uuid: UUID,
+ client: iperf_client.IPerfClientBase,
+ server_ip: str,
+ server_port: int,
+ active_streams: DictProxy[Any, Any],
+ stream_results: DictProxy[Any, Any],
+ access_category: Optional[str] = None,
+ bandwidth: Optional[int] = None,
+ stream_time: int = DEFAULT_STREAM_TIME,
+ start_time: Optional[float] = None,
):
"""Runs an iperf3 stream.
@@ -415,28 +417,21 @@
4. Removes stream UUID from active_streams
Args:
- uuid: UUID object, identifier for stream
+ uuid: Identifier for stream
client: IPerfClient object on device
- server_ip: string, ip address of IPerfServer for stream
- server_port: int, port of the IPerfServer for stream
- active_streams: multiprocessing.Manager.dict, which holds stream
- UUIDs of active streams on the device
- stream_results: multiprocessing.Manager.dict, which maps stream
- UUIDs of streams to IPerfResult objects
- access_category: string, WMM access category to use with iperf
- (AC_BK, AC_BE, AC_VI, AC_VO). Unset if None.
- bandwidth: int, bandwidth in mbps to use with iperf. Implies UDP.
- Unlimited if None.
- stream_time: int, time in seconds, to run iperf stream
- start_time: float, time, seconds since epoch, at which to start the
- stream (for better synchronicity). If None, start immediately.
+ server_ip: IP address of IPerfServer for stream
+ server_port: port of the IPerfServer for stream
+ active_streams: holds stream UUIDs of active streams on the device
+ stream_results: maps stream UUIDs of streams to IPerfResult objects
+ access_category: WMM access category to use with iperf (AC_BK, AC_BE, AC_VI,
+ AC_VO). Unset if None.
+ bandwidth: Bandwidth in mbps to use with iperf. Implies UDP. Unlimited if
+ None.
+ stream_time: Time in seconds, to run iperf stream
+ start_time: Time, seconds since epoch, at which to start the stream (for
+ better synchronicity). If None, start immediately.
"""
active_streams[uuid] = True
- # SSH sessions must be started within the process that is going to
- # use it.
- if type(client) == iperf_client.IPerfClientOverSsh:
- with utils.SuppressLogOutput():
- client.start_ssh()
ac_flag = ""
bandwidth_flag = ""
