]>
Commit | Line | Data |
---|---|---|
f67539c2 TL |
1 | #!/usr/bin/env python |
2 | ||
3 | import argparse | |
4 | import socket | |
5 | import sys | |
6 | ||
7 | from util import add_common_args | |
8 | from local_thrift import thrift # noqa | |
9 | from thrift.Thrift import TMessageType, TType | |
10 | from thrift.transport.TSocket import TSocket | |
11 | from thrift.transport.TTransport import TBufferedTransport, TFramedTransport | |
12 | from thrift.protocol.TBinaryProtocol import TBinaryProtocol | |
13 | from thrift.protocol.TCompactProtocol import TCompactProtocol | |
14 | ||
15 | ||
16 | def test_void(proto): | |
17 | proto.writeMessageBegin('testVoid', TMessageType.CALL, 3) | |
18 | proto.writeStructBegin('testVoid_args') | |
19 | proto.writeFieldStop() | |
20 | proto.writeStructEnd() | |
21 | proto.writeMessageEnd() | |
22 | proto.trans.flush() | |
23 | ||
24 | _, mtype, _ = proto.readMessageBegin() | |
25 | assert mtype == TMessageType.REPLY | |
26 | proto.readStructBegin() | |
27 | _, ftype, _ = proto.readFieldBegin() | |
28 | assert ftype == TType.STOP | |
29 | proto.readStructEnd() | |
30 | proto.readMessageEnd() | |
31 | ||
32 | ||
33 | # THeader stack should accept binary protocol with optionally framed transport | |
34 | def main(argv): | |
35 | p = argparse.ArgumentParser() | |
36 | add_common_args(p) | |
37 | # Since THeaderTransport acts as framed transport when detected frame, we | |
38 | # cannot use --transport=framed as it would result in 2 layered frames. | |
39 | p.add_argument('--override-transport') | |
40 | p.add_argument('--override-protocol') | |
41 | args = p.parse_args() | |
42 | assert args.protocol == 'header' | |
43 | assert args.transport == 'buffered' | |
44 | assert not args.ssl | |
45 | ||
46 | sock = TSocket(args.host, args.port, socket_family=socket.AF_INET) | |
47 | if not args.override_transport or args.override_transport == 'buffered': | |
48 | trans = TBufferedTransport(sock) | |
49 | elif args.override_transport == 'framed': | |
50 | print('TFRAMED') | |
51 | trans = TFramedTransport(sock) | |
52 | else: | |
53 | raise ValueError('invalid transport') | |
54 | trans.open() | |
55 | ||
56 | if not args.override_protocol or args.override_protocol == 'binary': | |
57 | proto = TBinaryProtocol(trans) | |
58 | elif args.override_protocol == 'compact': | |
59 | proto = TCompactProtocol(trans) | |
60 | else: | |
61 | raise ValueError('invalid transport') | |
62 | ||
63 | test_void(proto) | |
64 | test_void(proto) | |
65 | ||
66 | trans.close() | |
67 | ||
68 | ||
69 | if __name__ == '__main__': | |
70 | sys.exit(main(sys.argv[1:])) |