Mark Slee | 57cc25e | 2007-02-28 21:43:54 +0000 | [diff] [blame] | 1 | #!/usr/bin/env python |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 2 | |
David Reiss | ea2cba8 | 2009-03-30 21:35:00 +0000 | [diff] [blame] | 3 | # |
| 4 | # Licensed to the Apache Software Foundation (ASF) under one |
| 5 | # or more contributor license agreements. See the NOTICE file |
| 6 | # distributed with this work for additional information |
| 7 | # regarding copyright ownership. The ASF licenses this file |
| 8 | # to you under the Apache License, Version 2.0 (the |
| 9 | # "License"); you may not use this file except in compliance |
| 10 | # with the License. You may obtain a copy of the License at |
| 11 | # |
| 12 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 13 | # |
| 14 | # Unless required by applicable law or agreed to in writing, |
| 15 | # software distributed under the License is distributed on an |
| 16 | # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 17 | # KIND, either express or implied. See the License for the |
| 18 | # specific language governing permissions and limitations |
| 19 | # under the License. |
| 20 | # |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 21 | from __future__ import division |
David Reiss | db893b6 | 2008-02-18 02:11:48 +0000 | [diff] [blame] | 22 | import sys, glob, time |
Mark Slee | 5299a95 | 2007-10-05 00:13:24 +0000 | [diff] [blame] | 23 | sys.path.insert(0, './gen-py') |
| 24 | sys.path.insert(0, glob.glob('../../lib/py/build/lib.*')[0]) |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 25 | from optparse import OptionParser |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 26 | |
Mark Slee | 57cc25e | 2007-02-28 21:43:54 +0000 | [diff] [blame] | 27 | from ThriftTest import ThriftTest |
| 28 | from ThriftTest.ttypes import * |
Mark Slee | d788b2e | 2006-09-07 01:26:35 +0000 | [diff] [blame] | 29 | from thrift.transport import TTransport |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 30 | from thrift.transport import TSocket |
| 31 | from thrift.protocol import TBinaryProtocol |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 32 | from thrift.protocol import TCompactProtocol |
David Reiss | f78ec2b | 2009-01-31 21:59:32 +0000 | [diff] [blame] | 33 | from thrift.server import TServer, TNonblockingServer, THttpServer |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 34 | |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 35 | parser = OptionParser() |
| 36 | parser.set_defaults(port=9090, verbose=1, proto='binary') |
| 37 | parser.add_option("--port", type="int", dest="port", |
| 38 | help="port number for server to listen on") |
| 39 | parser.add_option('-v', '--verbose', action="store_const", |
| 40 | dest="verbose", const=2, |
| 41 | help="verbose output") |
| 42 | parser.add_option('--proto', dest="proto", type="string", |
| 43 | help="protocol to use, one of: accel, binary, compact") |
| 44 | options, args = parser.parse_args() |
| 45 | |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 46 | class TestHandler: |
| 47 | |
| 48 | def testVoid(self): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 49 | if options.verbose: |
| 50 | print 'testVoid()' |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 51 | |
| 52 | def testString(self, str): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 53 | if options.verbose: |
| 54 | print 'testString(%s)' % str |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 55 | return str |
| 56 | |
| 57 | def testByte(self, byte): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 58 | if options.verbose: |
| 59 | print 'testByte(%d)' % byte |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 60 | return byte |
| 61 | |
Mark Slee | c98d050 | 2006-09-06 02:42:25 +0000 | [diff] [blame] | 62 | def testI16(self, i16): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 63 | if options.verbose: |
| 64 | print 'testI16(%d)' % i16 |
Mark Slee | c98d050 | 2006-09-06 02:42:25 +0000 | [diff] [blame] | 65 | return i16 |
| 66 | |
| 67 | def testI32(self, i32): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 68 | if options.verbose: |
| 69 | print 'testI32(%d)' % i32 |
Mark Slee | c98d050 | 2006-09-06 02:42:25 +0000 | [diff] [blame] | 70 | return i32 |
| 71 | |
| 72 | def testI64(self, i64): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 73 | if options.verbose: |
| 74 | print 'testI64(%d)' % i64 |
Mark Slee | c98d050 | 2006-09-06 02:42:25 +0000 | [diff] [blame] | 75 | return i64 |
| 76 | |
| 77 | def testDouble(self, dub): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 78 | if options.verbose: |
| 79 | print 'testDouble(%f)' % dub |
Mark Slee | c98d050 | 2006-09-06 02:42:25 +0000 | [diff] [blame] | 80 | return dub |
| 81 | |
| 82 | def testStruct(self, thing): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 83 | if options.verbose: |
| 84 | print 'testStruct({%s, %d, %d, %d})' % (thing.string_thing, thing.byte_thing, thing.i32_thing, thing.i64_thing) |
Mark Slee | c98d050 | 2006-09-06 02:42:25 +0000 | [diff] [blame] | 85 | return thing |
| 86 | |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 87 | def testException(self, str): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 88 | if options.verbose: |
| 89 | print 'testException(%s)' % str |
Mark Slee | c98d050 | 2006-09-06 02:42:25 +0000 | [diff] [blame] | 90 | if str == 'Xception': |
| 91 | x = Xception() |
| 92 | x.errorCode = 1001 |
| 93 | x.message = str |
| 94 | raise x |
David Reiss | bcaa2ad | 2008-06-10 22:55:26 +0000 | [diff] [blame] | 95 | elif str == "throw_undeclared": |
| 96 | raise ValueError("foo") |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 97 | |
David Reiss | 6ce401d | 2009-03-24 20:01:58 +0000 | [diff] [blame] | 98 | def testOneway(self, seconds): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 99 | if options.verbose: |
| 100 | print 'testOneway(%d) => sleeping...' % seconds |
| 101 | time.sleep(seconds / 3) # be quick |
| 102 | if options.verbose: |
| 103 | print 'done sleeping' |
David Reiss | db893b6 | 2008-02-18 02:11:48 +0000 | [diff] [blame] | 104 | |
David Reiss | 7442127 | 2008-11-07 23:09:31 +0000 | [diff] [blame] | 105 | def testNest(self, thing): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 106 | if options.verbose: |
| 107 | print 'testNest(%s)' % thing |
David Reiss | 7442127 | 2008-11-07 23:09:31 +0000 | [diff] [blame] | 108 | return thing |
| 109 | |
| 110 | def testMap(self, thing): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 111 | if options.verbose: |
| 112 | print 'testMap(%s)' % thing |
David Reiss | 7442127 | 2008-11-07 23:09:31 +0000 | [diff] [blame] | 113 | return thing |
| 114 | |
| 115 | def testSet(self, thing): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 116 | if options.verbose: |
| 117 | print 'testSet(%s)' % thing |
David Reiss | 7442127 | 2008-11-07 23:09:31 +0000 | [diff] [blame] | 118 | return thing |
| 119 | |
| 120 | def testList(self, thing): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 121 | if options.verbose: |
| 122 | print 'testList(%s)' % thing |
David Reiss | 7442127 | 2008-11-07 23:09:31 +0000 | [diff] [blame] | 123 | return thing |
| 124 | |
| 125 | def testEnum(self, thing): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 126 | if options.verbose: |
| 127 | print 'testEnum(%s)' % thing |
David Reiss | 7442127 | 2008-11-07 23:09:31 +0000 | [diff] [blame] | 128 | return thing |
| 129 | |
| 130 | def testTypedef(self, thing): |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 131 | if options.verbose: |
| 132 | print 'testTypedef(%s)' % thing |
David Reiss | 7442127 | 2008-11-07 23:09:31 +0000 | [diff] [blame] | 133 | return thing |
| 134 | |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 135 | def testMapMap(self, thing): |
| 136 | if options.verbose: |
| 137 | print 'testMapMap(%s)' % thing |
| 138 | return thing |
| 139 | |
| 140 | def testMulti(self, arg0, arg1, arg2, arg3, arg4, arg5): |
| 141 | if options.verbose: |
| 142 | print 'testMulti(%s)' % [arg0, arg1, arg2, arg3, arg4, arg5] |
| 143 | x = Xtruct(byte_thing=arg0, i32_thing=arg1, i64_thing=arg2) |
| 144 | return x |
| 145 | |
| 146 | if options.proto == 'binary': |
| 147 | pfactory = TBinaryProtocol.TBinaryProtocolFactory() |
| 148 | elif options.proto == 'accel': |
| 149 | pfactory = TBinaryProtocol.TBinaryProtocolAcceleratedFactory() |
| 150 | elif options.proto == 'compact': |
| 151 | pfactory = TCompactProtocol.TCompactProtocolFactory() |
| 152 | else: |
| 153 | raise AssertionError('Unknown --proto option: %s' % options.proto) |
Mark Slee | c967656 | 2006-09-05 17:34:52 +0000 | [diff] [blame] | 154 | handler = TestHandler() |
Mark Slee | 1dd819c | 2006-10-26 04:56:18 +0000 | [diff] [blame] | 155 | processor = ThriftTest.Processor(handler) |
David Reiss | bcaa2ad | 2008-06-10 22:55:26 +0000 | [diff] [blame] | 156 | |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 157 | if args[0] == "THttpServer": |
| 158 | server = THttpServer.THttpServer(processor, ('', options.port), pfactory) |
David Reiss | 7442127 | 2008-11-07 23:09:31 +0000 | [diff] [blame] | 159 | else: |
Roger Meier | 4ebaa76 | 2011-02-08 20:44:22 +0000 | [diff] [blame] | 160 | host = None |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 161 | transport = TSocket.TServerSocket(host, options.port) |
David Reiss | f78ec2b | 2009-01-31 21:59:32 +0000 | [diff] [blame] | 162 | tfactory = TTransport.TBufferedTransportFactory() |
| 163 | |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 164 | if args[0] == "TNonblockingServer": |
| 165 | server = TNonblockingServer.TNonblockingServer(processor, transport, inputProtocolFactory=pfactory) |
| 166 | elif args[0] == "TProcessPoolServer": |
| 167 | import signal |
| 168 | def set_alarm(): |
| 169 | def clean_shutdown(signum, frame): |
| 170 | for worker in server.workers: |
| 171 | print 'Terminating worker: %s' % worker |
| 172 | worker.terminate() |
| 173 | print 'Requesting server to stop()' |
| 174 | server.stop() |
| 175 | signal.signal(signal.SIGALRM, clean_shutdown) |
| 176 | signal.alarm(2) |
| 177 | from thrift.server import TProcessPoolServer |
| 178 | server = TProcessPoolServer.TProcessPoolServer(processor, transport, tfactory, pfactory) |
| 179 | server.setNumWorkers(5) |
| 180 | set_alarm() |
David Reiss | f78ec2b | 2009-01-31 21:59:32 +0000 | [diff] [blame] | 181 | else: |
Bryan Duxbury | 59d4efd | 2011-03-21 17:38:22 +0000 | [diff] [blame^] | 182 | ServerClass = getattr(TServer, args[0]) |
David Reiss | f78ec2b | 2009-01-31 21:59:32 +0000 | [diff] [blame] | 183 | server = ServerClass(processor, transport, tfactory, pfactory) |
| 184 | |
Mark Slee | 794993d | 2006-09-20 01:56:10 +0000 | [diff] [blame] | 185 | server.serve() |