NDOJ/judge/bridge/echo_test_client.py

import os
import socket
import struct
import time
import zlib

size_pack = struct.Struct("!I")


def open_connection():
    sock = socket.create_connection((host, port))
    return sock


def zlibify(data):
    data = zlib.compress(data.encode("utf-8"))
    return size_pack.pack(len(data)) + data


def dezlibify(data, skip_head=True):
    if skip_head:
        data = data[size_pack.size :]
    return zlib.decompress(data).decode("utf-8")


def main():
    global host, port
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("-l", "--host", default="localhost")
    parser.add_argument("-p", "--port", default=9999, type=int)
    args = parser.parse_args()
    host, port = args.host, args.port

    print("Opening idle connection:", end=" ")
    s1 = open_connection()
    print("Success")
    print("Opening hello world connection:", end=" ")
    s2 = open_connection()
    print("Success")
    print("Sending Hello, World!", end=" ")
    s2.sendall(zlibify("Hello, World!"))
    print("Success")
    print("Testing blank connection:", end=" ")
    s3 = open_connection()
    s3.close()
    print("Success")
    result = dezlibify(s2.recv(1024))
    assert result == "Hello, World!"
    print(result)
    s2.close()
    print("Large random data test:", end=" ")
    s4 = open_connection()
    data = os.urandom(1000000).decode("iso-8859-1")
    print("Generated", end=" ")
    s4.sendall(zlibify(data))
    print("Sent", end=" ")
    result = b""
    while len(result) < size_pack.size:
        result += s4.recv(1024)
    size = size_pack.unpack(result[: size_pack.size])[0]
    result = result[size_pack.size :]
    while len(result) < size:
        result += s4.recv(1024)
    print("Received", end=" ")
    assert dezlibify(result, False) == data
    print("Success")
    s4.close()
    print("Test malformed connection:", end=" ")
    s5 = open_connection()
    s5.sendall(data[:100000].encode("utf-8"))
    s5.close()
    print("Success")
    print("Waiting for timeout to close idle connection:", end=" ")
    time.sleep(6)
    print("Done")
    s1.close()


if __name__ == "__main__":
    main()