1#!/usr/bin/env python
2
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#
21
22from thrift.transport import TSocket
23import unittest
24import time
25import socket
26import random
27
28
29class TimeoutTest(unittest.TestCase):
30    def setUp(self):
31        for i in range(50):
32            try:
33                # find a port we can use
34                self.listen_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
35                self.port = random.randint(10000, 30000)
36                self.listen_sock.bind(('localhost', self.port))
37                self.listen_sock.listen(5)
38                break
39            except Exception:
40                if i == 49:
41                    raise
42
43    def testConnectTimeout(self):
44        starttime = time.time()
45
46        try:
47            leaky = []
48            for i in range(100):
49                socket = TSocket.TSocket('localhost', self.port)
50                socket.setTimeout(10)
51                socket.open()
52                leaky.append(socket)
53        except Exception:
54            self.assert_(time.time() - starttime < 5.0)
55
56    def testWriteTimeout(self):
57        starttime = time.time()
58
59        try:
60            socket = TSocket.TSocket('localhost', self.port)
61            socket.setTimeout(10)
62            socket.open()
63            lsock = self.listen_sock.accept()
64            while True:
65                lsock.write("hi" * 100)
66
67        except Exception:
68            self.assert_(time.time() - starttime < 5.0)
69
70
71if __name__ == '__main__':
72    suite = unittest.TestSuite()
73    loader = unittest.TestLoader()
74
75    suite.addTest(loader.loadTestsFromTestCase(TimeoutTest))
76
77    testRunner = unittest.TextTestRunner(verbosity=2)
78    testRunner.run(suite)
79