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