pytact.fake_coq_client

 1import contextlib
 2import socket
 3import argparse
 4import signal
 5from subprocess import Popen
 6import capnp
 7import pytact.graph_api_capnp as graph_api_capnp
 8from pytact.data_reader import capnp_message_generator_from_file_lowlevel, record_lowlevel_generator
 9
10def run_fake_client(server_socket, messages_generator):
11    socket_reader = graph_api_capnp.PredictionProtocol.Response.read_multiple_packed(
12        server_socket, traversal_limit_in_words=2**64-1)
13    for msg in messages_generator:
14        msg.dynamic.as_builder().write_packed(server_socket)
15        prev_sig = signal.signal(signal.SIGINT, signal.SIG_DFL)  # SIGINT catching OFF
16        response = next(socket_reader)
17        signal.signal(signal.SIGINT, prev_sig)  # SIGINT catching ON
18        messages_generator.send(response.as_builder())
19
20def compare(request, response, recorded_response):
21    if response.to_dict() == recorded_response.to_dict():
22        print(f'The servers response to a {request.which.name} message was equal to the '
23              f'recorded response')
24    else:
25        raise ValueError(
26            f"The servers response to a {request.which.name} message was not equal to the " +
27            f"recorded response.\n"
28            f"Recorded response: {recorded_response}\n"
29            f"Servers response: {response}\n"
30        )
31
32def main():
33    parser = argparse.ArgumentParser(
34        description = 'A fake Coq client that connects to a prediction server and feeds it a stream of previously ' +
35                      'recorded messages.',
36        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
37    parser.add_argument('data',
38                        type=str,
39                        help='A file that contains a previously recorded communication sequence between Coq ' +
40                             'and a prediction server.')
41    parser.add_argument('--check',
42                        action=argparse.BooleanOptionalAction,
43                        default = True,
44                        help='Wether or not to compare the response messages of the server to recorded messages.')
45    parser.add_argument('--record',
46                        dest="record_file",
47                        type = str,
48                        default = None,
49                        help='Re-record the interaction with the server to a new file (useful when you want to ' +
50                    'update the trace with a new response behavior of a server)')
51    connect_group = parser.add_mutually_exclusive_group(required=True)
52    connect_group.add_argument('--tcp',
53                        dest='tcp_location',
54                        type = str,
55                        default = None,
56                        help='Connect to a prediction server on the specified ip:port.')
57    connect_group.add_argument('--stdin',
58                        dest='executable',
59                        type = str,
60                        default = None,
61                        help='Start the specified prediction server and connect to it through stdin.')
62    cmd_args = parser.parse_args()
63
64    with open(cmd_args.data, 'rb') as message_file:
65        messages_generator = capnp_message_generator_from_file_lowlevel(
66            message_file,
67            check= compare if cmd_args.check else None)
68        if cmd_args.record_file is not None:
69            record_context = open(cmd_args.record_file, 'wb')
70        else:
71            record_context = contextlib.nullcontext()
72        with record_context as record_file:
73            if record_file is not None:
74                messages_generator = record_lowlevel_generator(record_file, messages_generator)
75            if cmd_args.tcp_location is not None:
76                addr, port = cmd_args.tcp_location.split(':')
77                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:
78                    server_socket.connect((addr, int(port)))
79                    run_fake_client(server_socket, messages_generator)
80            else:
81                our_sock, their_sock = socket.socketpair()
82                process = Popen(cmd_args.executable, shell=True, stdin=their_sock)
83                their_sock.close()
84                run_fake_client(our_sock, messages_generator)
85                our_sock.close()
86                process.communicate()
87
88if __name__ == '__main__':
89    main()