ESP32/control_iperf.py

123 lines
4.6 KiB
Python
Executable File

#!/usr/bin/env python3
import asyncio
import argparse
import serial_asyncio
import sys
import re
class SerialController(asyncio.Protocol):
def __init__(self, port_name, args, loop, completion_future):
self.port_name = port_name
self.args = args
self.loop = loop
self.transport = None
self.buffer = ""
self.completion_future = completion_future
if args.action == 'pps':
self.cmd_str = f"iperf pps {args.value}\n"
self.target_key = "IPERF_PPS_UPDATED"
elif args.action == 'status':
self.cmd_str = "iperf status\n"
self.target_key = "IPERF_STATUS"
elif args.action == 'start':
self.cmd_str = "iperf start\n"
self.target_key = "IPERF_STARTED"
elif args.action == 'stop':
self.cmd_str = "iperf stop\n"
self.target_key = "IPERF_STOPPED"
def connection_made(self, transport):
self.transport = transport
transport.write(b'\n')
self.loop.create_task(self.send_command())
async def send_command(self):
await asyncio.sleep(0.1)
self.transport.write(self.cmd_str.encode())
def data_received(self, data):
self.buffer += data.decode(errors='ignore')
# FIX: Process complete lines only to avoid partial regex matching
while '\n' in self.buffer:
line, self.buffer = self.buffer.split('\n', 1)
line = line.strip()
if self.target_key in line:
if not self.completion_future.done():
if self.args.action == 'status':
# FIX: Added [-]? to allow negative error rates (overshoot)
# Regex: Err=([-\d\.]+)%
m = re.search(r'Running=(\d+), Config=(\d+), Actual=(\d+), Err=([-\d\.]+)%', line)
if m:
state = "Running" if m.group(1) == '1' else "Stopped"
self.completion_future.set_result(f"{state}, Cfg: {m.group(2)}, Act: {m.group(3)}, Err: {m.group(4)}%")
else:
# Now if it fails, it's a true format mismatch, not fragmentation
self.completion_future.set_result(f"Parse Error on line: {line}")
else:
self.completion_future.set_result(True)
self.transport.close()
return
def connection_lost(self, exc):
if not self.completion_future.done():
self.completion_future.set_exception(Exception("Closed"))
async def run_device(port, args):
loop = asyncio.get_running_loop()
fut = loop.create_future()
try:
await serial_asyncio.create_serial_connection(
loop, lambda: SerialController(port, args, loop, fut), port, baudrate=115200)
return await asyncio.wait_for(fut, timeout=2.0)
except:
return None
def expand_devices(device_str):
devices = []
parts = [d.strip() for d in device_str.split(',')]
for part in parts:
match = re.match(r'^(.*?)(\d+)-(\d+)$', part)
if match:
prefix, start, end = match.group(1), int(match.group(2)), int(match.group(3))
step = 1 if end >= start else -1
for i in range(start, end + step, step):
devices.append(f"{prefix}{i}")
else:
devices.append(part)
return devices
async def main():
parser = argparse.ArgumentParser()
parser.add_argument('action', choices=['start', 'stop', 'pps', 'status'])
parser.add_argument('value_arg', nargs='?', type=int, help='Value for PPS')
parser.add_argument('--value', type=int, help='Value for PPS')
parser.add_argument('--devices', required=True, help="/dev/ttyUSB0-29")
args = parser.parse_args()
if args.value_arg is not None: args.value = args.value_arg
if args.action == 'pps' and args.value is None:
print("Error: 'pps' action requires a value")
sys.exit(1)
if sys.platform == 'win32': asyncio.set_event_loop(asyncio.ProactorEventLoop())
devs = expand_devices(args.devices)
print(f"Executing '{args.action}' on {len(devs)} devices...")
tasks = [run_device(d, args) for d in devs]
results = await asyncio.gather(*tasks)
print("\nResults:")
for dev, res in zip(devs, results):
if args.action == 'status':
print(f"{dev}: {res if res else 'TIMEOUT'}")
else:
status = "OK" if res is True else "FAIL"
print(f"{dev}: {status}")
if __name__ == '__main__':
asyncio.run(main())