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_string(proto, value):
13    method_name = 'testString'
14    ttype = TType.STRING
15    proto.writeMessageBegin(method_name, TMessageType.CALL, 3)
16    proto.writeStructBegin(method_name + '_args')
17    proto.writeFieldBegin('thing', ttype, 1)
18    proto.writeString(value)
19    proto.writeFieldEnd()
20    proto.writeFieldStop()
21    proto.writeStructEnd()
22    proto.writeMessageEnd()
23    proto.trans.flush()
24
25    _, mtype, _ = proto.readMessageBegin()
26    assert mtype == TMessageType.REPLY
27    proto.readStructBegin()
28    _, ftype, fid = proto.readFieldBegin()
29    assert fid == 0
30    assert ftype == ttype
31    result = proto.readString()
32    proto.readFieldEnd()
33    _, ftype, _ = proto.readFieldBegin()
34    assert ftype == TType.STOP
35    proto.readStructEnd()
36    proto.readMessageEnd()
37    assert value == result
38
39
40def main(argv):
41    p = argparse.ArgumentParser()
42    add_common_args(p)
43    p.add_argument('--limit', type=int)
44    args = p.parse_args()
45    proto = init_protocol(args)
46    test_string(proto, 'a' * (args.limit - 1))
47    test_string(proto, 'a' * (args.limit - 1))
48    print('[OK]: limit - 1')
49    test_string(proto, 'a' * args.limit)
50    test_string(proto, 'a' * args.limit)
51    print('[OK]: just limit')
52    try:
53        test_string(proto, 'a' * (args.limit + 1))
54    except Exception:
55        print('[OK]: limit + 1')
56    else:
57        print('[ERROR]: limit + 1')
58        assert False
59
60
61if __name__ == '__main__':
62    main(sys.argv[1:])
63