1#!/usr/bin/env python
2
3import argparse
4import sys
5
6from util import add_common_args, init_protocol
7from local_thrift import thrift  # noqa
8from thrift.Thrift import TMessageType, TType
9
10
11# TODO: generate from ThriftTest.thrift
12def test_list(proto, value):
13    method_name = 'testList'
14    ttype = TType.LIST
15    etype = TType.I32
16    proto.writeMessageBegin(method_name, TMessageType.CALL, 3)
17    proto.writeStructBegin(method_name + '_args')
18    proto.writeFieldBegin('thing', ttype, 1)
19    proto.writeListBegin(etype, len(value))
20    for e in value:
21        proto.writeI32(e)
22    proto.writeListEnd()
23    proto.writeFieldEnd()
24    proto.writeFieldStop()
25    proto.writeStructEnd()
26    proto.writeMessageEnd()
27    proto.trans.flush()
28
29    _, mtype, _ = proto.readMessageBegin()
30    assert mtype == TMessageType.REPLY
31    proto.readStructBegin()
32    _, ftype, fid = proto.readFieldBegin()
33    assert fid == 0
34    assert ftype == ttype
35    etype2, len2 = proto.readListBegin()
36    assert etype == etype2
37    assert len2 == len(value)
38    for i in range(len2):
39        v = proto.readI32()
40        assert v == value[i]
41    proto.readListEnd()
42    proto.readFieldEnd()
43    _, ftype, _ = proto.readFieldBegin()
44    assert ftype == TType.STOP
45    proto.readStructEnd()
46    proto.readMessageEnd()
47
48
49def main(argv):
50    p = argparse.ArgumentParser()
51    add_common_args(p)
52    p.add_argument('--limit', type=int)
53    args = p.parse_args()
54    proto = init_protocol(args)
55    # TODO: test set and map
56    test_list(proto, list(range(args.limit - 1)))
57    test_list(proto, list(range(args.limit - 1)))
58    print('[OK]: limit - 1')
59    test_list(proto, list(range(args.limit)))
60    test_list(proto, list(range(args.limit)))
61    print('[OK]: just limit')
62    try:
63        test_list(proto, list(range(args.limit + 1)))
64    except Exception:
65        print('[OK]: limit + 1')
66    else:
67        print('[ERROR]: limit + 1')
68        assert False
69
70
71if __name__ == '__main__':
72    sys.exit(main(sys.argv[1:]))
73