Update bridge (DMOJ)

judge/bridge/echo_test_client.py

@@ -0,0 +1,81 @@
+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()