wangyu-UDPspeeder/test/udp_speeder_test.py

import socket
import subprocess
import sys
import threading
import signal
import time
import argparse

LOG_LEVEL_MAP = {
    "never": 0,
    "fatal": 1,
    "error": 2,
    "warn": 3,
    "info": 4,
    "debug": 5,
    "trace": 6
}

class UDPspeederTest:
    def __init__(self, speederv2_path, log_level="info", socket_activate=False, extra_udpspeeder_args="", no_udpspeeder=False):
        self.speederv2_path = speederv2_path
        self.log_level = log_level
        self.socket_activate = socket_activate
        self.extra_udpspeeder_args = extra_udpspeeder_args
        self.no_udpspeeder = no_udpspeeder

    def run_speederv2(self, mode, local_port, remote_ip, remote_port):
        log_level_num = LOG_LEVEL_MAP.get(self.log_level, 4)  # Default to "info" if log_level is not found
        cmd = f"{self.speederv2_path} {mode} -l0.0.0.0:{local_port} -r{remote_ip}:{remote_port} -f20:10 --log-level {log_level_num} {self.extra_udpspeeder_args}"
        if self.socket_activate:
            cmd = f"systemd-socket-activate -l0.0.0.0:{local_port} -d {cmd}"
        print(f"UPDspeeder command: {cmd}")
        
        process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        return process

    def start_udpspeeder_server(self, local_port, remote_port):
        return self.run_speederv2("-s", local_port, "127.0.0.1", remote_port)

    def start_udpspeeder_client(self, local_port, remote_ip, remote_port):
        return self.run_speederv2("-c", local_port, remote_ip, remote_port)

    def server(self, udp_speeder_port, server_port):
        # Start UDPspeeder server if not disabled
        if not self.no_udpspeeder:
            udpspeeder_process = self.start_udpspeeder_server(udp_speeder_port, server_port)
            threading.Thread(target=self.log_udpspeeder_output, args=(udpspeeder_process,)).start()

        # Start UDP server
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.bind(("0.0.0.0", server_port))
        print(f"Server listening on port {server_port}")

        def cleanup(signum, frame):
            print("Shutting down server...")
            sock.close()
            if not self.no_udpspeeder:
                udpspeeder_process.terminate()
            sys.exit(0)

        signal.signal(signal.SIGTERM, cleanup)

        while True:
            data, addr = sock.recvfrom(1024)
            print(f"Received message: {data.decode()} from {addr}")
            reply = f"Server reply to {data.decode()}"
            sock.sendto(reply.encode(), addr)
            print(f"Sent reply: {reply} ==============================================================")

    def client(self, us_client_port, us_server_port):
        # Start UDPspeeder client if not disabled
        if not self.no_udpspeeder:
            udpspeeder_process = self.start_udpspeeder_client(us_client_port, "127.0.0.1", us_server_port)
            threading.Thread(target=self.log_udpspeeder_output, args=(udpspeeder_process,)).start()

        # Start UDP client
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        server_address = ("127.0.0.1", us_client_port)
        message = "Hello, Server!"

        def cleanup(signum, frame):
            print("Shutting down client...")
            sock.close()
            if not self.no_udpspeeder:
                udpspeeder_process.terminate()
            sys.exit(0)

        signal.signal(signal.SIGTERM, cleanup)
        sock.settimeout(1.0)

        while True:
            print(f"Sending message: {message}")
            sent = sock.sendto(message.encode(), server_address)

            try:
                data, server = sock.recvfrom(1024)
            except socket.timeout:
                print("Request timed out")
                continue

            print(f"Received reply: {data.decode()} ==============================================================")
            time.sleep(1)

        sock.close()

    def log_udpspeeder_output(self, process):
        for line in process.stdout:
            print(line.decode(), end='')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="UDPspeeder test script")
    parser.add_argument("--mode", choices=["server", "client"], help="Mode to run: 'server' or 'client'")
    parser.add_argument("--us-server-port", type=int, default=4096, help="UDPspeeder server port")
    parser.add_argument("--us-client-port", type=int, default=3333, help="UDPspeeder client port")
    parser.add_argument("--server-port", type=int, default=7777, help="Server port behind UDPspeeder")
    parser.add_argument("--log-level", choices=LOG_LEVEL_MAP.keys(), default="info", nargs="?", help="Log level: 'never', 'fatal', 'error', 'warn', 'info', 'debug', 'trace'")
    parser.add_argument("--socket-activate", action="store_true", help="Test systemd socket activation")
    parser.add_argument("--extra-args", help="Extra arguments to pass to UDPspeeder")
    parser.add_argument("--no-udpspeeder", action="store_true", help="Do not start UDPspeeder")
    parser.add_argument("--speederv2-path", default="speederv2", help="Path to the speederv2 binary")


    args = parser.parse_args()

    udpspeeder_test = UDPspeederTest(args.speederv2_path, args.log_level, args.socket_activate, args.extra_args, args.no_udpspeeder)

    if args.mode == "server":
        udpspeeder_test.server(args.us_server_port, args.server_port)
    elif args.mode == "client":
        udpspeeder_test.client(args.us_client_port, args.us_server_port)
    else:
        print("Invalid mode. Use 'server' or 'client'.")
        sys.exit(1)