ceph/src/jaegertracing/thrift/test/features/util.py

   1 import argparse
   2 import socket
   3
   4 from local_thrift import thrift  # noqa
   5 from thrift.transport.TSocket import TSocket
   6 from thrift.transport.TTransport import TBufferedTransport, TFramedTransport
   7 from thrift.transport.THttpClient import THttpClient
   8 from thrift.protocol.TBinaryProtocol import TBinaryProtocol
   9 from thrift.protocol.TCompactProtocol import TCompactProtocol
  10 from thrift.protocol.TJSONProtocol import TJSONProtocol
  11
  12
  13 def add_common_args(p):
  14     p.add_argument('--host', default='localhost')
  15     p.add_argument('--port', type=int, default=9090)
  16     p.add_argument('--protocol', default='binary')
  17     p.add_argument('--transport', default='buffered')
  18     p.add_argument('--ssl', action='store_true')
  19
  20
  21 def parse_common_args(argv):
  22     p = argparse.ArgumentParser()
  23     add_common_args(p)
  24     return p.parse_args(argv)
  25
  26
  27 def init_protocol(args):
  28     sock = TSocket(args.host, args.port, socket_family=socket.AF_INET)
  29     sock.setTimeout(500)
  30     trans = {
  31         'buffered': TBufferedTransport,
  32         'framed': TFramedTransport,
  33         'http': THttpClient,
  34     }[args.transport](sock)
  35     trans.open()
  36     return {
  37         'binary': TBinaryProtocol,
  38         'compact': TCompactProtocol,
  39         'json': TJSONProtocol,
  40     }[args.protocol](trans)