main.py

import argparse
import logging
import os

from iperf import get_iperf_server
from iperf import iperf_client
import iperf
import plot_results_1
import plot_results_2
import convert_to_csv

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('iperf_client.log', mode='w'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

def main():
    parser = argparse.ArgumentParser(description='iPerf client')
    parser.add_argument('-n', '--number', help='Number of servers to connect to', type=int, default=10)
    parser.add_argument('-np', '--numberplot', help='Number of servers to plot', type=int, default=1)
    parser.add_argument('-l', '--len', help='Block size in bytes (default: 128KB)', type=int, default=128*1024)
    parser.add_argument('-d', '--duration', help='Test duration in seconds (default: 10)', type=int, default=10)
    parser.add_argument('--num_plot', help='Number of servers to plot (default: all)', type=int, default=1)
    parser.add_argument('-p', '--parallel', help='Number of parallel streams (not implemented) (default: 1)', type=int, default=1)
    parser.add_argument('--goodput_location', help='Where goodput json is stored', type=str, default="results/goodput.json")
    parser.add_argument('--marker_string', help='Marker string to identify runs in logs/plots', type=str, default="")
    parser.add_argument('--debug', help='Enable debug logging', action='store_true')
    args = parser.parse_args()

    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)

    
    if os.path.exists(args.goodput_location):
        os.remove(args.goodput_location)
        logger.info(f"Cleared existing results file: {args.goodput_location}")
    
    servers = get_iperf_server.get_random_servers(args.number)
    setting = iperf.iperf_setting.iperfSetting(
        serv=None,
        blksize=args.len,
        duration=args.duration,
        parallel=args.parallel,
    )
    
    for serv in servers:
        setting.serv = serv
        code = iperf_client.run_iperf(setting, output_path=args.goodput_location)
        if code == -1:
            logger.info("Error Occurred, will pick a new server")
            new_server = get_iperf_server.get_random_servers(1)[0]
            while new_server in servers:
                new_server = get_iperf_server.get_random_servers(1)[0]
            servers.append(new_server)
    if args.marker_string == "":
        plot_results_1.plot_goodput(args.goodput_location, output_path=f"results/goodput_plot.pdf")
        plot_results_2.plot_tcpstats(args.goodput_location, output_path=f"results/tcpstats_plot.pdf", n=args.numberplot)
    else:
        plot_results_1.plot_goodput(args.goodput_location, output_path=f"results/{args.marker_string}_goodput_plot.pdf")
        plot_results_2.plot_tcpstats(args.goodput_location, output_path=f"results/{args.marker_string}_tcpstats_plot.pdf", n=args.numberplot)
    
    convert_to_csv.convert_to_csv(args.goodput_location, marker=args.marker_string)

if __name__ == '__main__':
    main()