← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/tests/test_test_server.py

Merge http-leaks into sftp-leaks

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/tests/test_test_server.py
 
1
# Copyright (C) 2010 Canonical Ltd
 
2
#
 
3
# This program is free software; you can redistribute it and/or modify
 
4
# it under the terms of the GNU General Public License as published by
 
5
# the Free Software Foundation; either version 2 of the License, or
 
6
# (at your option) any later version.
 
7
#
 
8
# This program is distributed in the hope that it will be useful,
 
9
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
10
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
11
# GNU General Public License for more details.
 
12
#
 
13
# You should have received a copy of the GNU General Public License
 
14
# along with this program; if not, write to the Free Software
 
15
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 
16
 
 
17
import errno
 
18
import socket
 
19
import SocketServer
 
20
 
 
21
from bzrlib import (
 
22
    osutils,
 
23
    tests,
 
24
    )
 
25
from bzrlib.tests import test_server
 
26
 
 
27
 
 
28
def load_tests(basic_tests, module, loader):
 
29
    suite = loader.suiteClass()
 
30
    server_tests, remaining_tests = tests.split_suite_by_condition(
 
31
        basic_tests, tests.condition_isinstance(TestTCPServerInAThread))
 
32
    server_scenarios = [ (name, {'server_class': getattr(test_server, name)})
 
33
                         for name in
 
34
                         ('TestingTCPServer', 'TestingThreadingTCPServer')]
 
35
    tests.multiply_tests(server_tests, server_scenarios, suite)
 
36
    suite.addTest(remaining_tests)
 
37
    return suite
 
38
 
 
39
 
 
40
class TCPClient(object):
 
41
 
 
42
    def __init__(self):
 
43
        self.sock = None
 
44
 
 
45
    def connect(self, addr):
 
46
        if self.sock is not None:
 
47
            raise AssertionError('Already connected to %r'
 
48
                                 % (self.sock.getsockname(),))
 
49
        self.sock = osutils.connect_socket(addr)
 
50
 
 
51
    def disconnect(self):
 
52
        if self.sock is not None:
 
53
            try:
 
54
                self.sock.shutdown(socket.SHUT_RDWR)
 
55
                self.sock.close()
 
56
            except socket.error, e:
 
57
                if e[0] in (errno.EBADF, errno.ENOTCONN):
 
58
                    # Right, the socket is already down
 
59
                    pass
 
60
                else:
 
61
                    raise
 
62
            self.sock = None
 
63
 
 
64
    def write(self, s):
 
65
        return self.sock.sendall(s)
 
66
 
 
67
    def read(self, bufsize=4096):
 
68
        return self.sock.recv(bufsize)
 
69
 
 
70
 
 
71
class TCPConnectionHandler(SocketServer.StreamRequestHandler):
 
72
 
 
73
    def handle(self):
 
74
        self.done = False
 
75
        self.handle_connection()
 
76
        while not self.done:
 
77
            self.handle_connection()
 
78
 
 
79
    def handle_connection(self):
 
80
        req = self.rfile.readline()
 
81
        if not req:
 
82
            self.done = True
 
83
        elif req == 'ping\n':
 
84
            self.wfile.write('pong\n')
 
85
        else:
 
86
            raise ValueError('[%s] not understood' % req)
 
87
 
 
88
 
 
89
class TestTCPServerInAThread(tests.TestCase):
 
90
 
 
91
    # Set by load_tests()
 
92
    server_class = None
 
93
 
 
94
    def get_server(self, server_class=None, connection_handler_class=None):
 
95
        if server_class is not None:
 
96
            self.server_class = server_class
 
97
        if connection_handler_class is None:
 
98
            connection_handler_class = TCPConnectionHandler
 
99
        server =  test_server.TestingTCPServerInAThread(
 
100
            ('localhost', 0), self.server_class, connection_handler_class)
 
101
        server.start_server()
 
102
        self.addCleanup(server.stop_server)
 
103
        return server
 
104
 
 
105
    def get_client(self):
 
106
        client = TCPClient()
 
107
        self.addCleanup(client.disconnect)
 
108
        return client
 
109
 
 
110
    def get_server_connection(self, server, conn_rank):
 
111
        return server.server.clients[conn_rank]
 
112
 
 
113
    def assertClientAddr(self, client, server, conn_rank):
 
114
        conn = self.get_server_connection(server, conn_rank)
 
115
        self.assertEquals(client.sock.getsockname(), conn[1])
 
116
 
 
117
    def test_start_stop(self):
 
118
        server = self.get_server()
 
119
        client = self.get_client()
 
120
        client.connect((server.host, server.port))
 
121
        server.stop_server()
 
122
        # since the server doesn't accept connections anymore attempting to
 
123
        # connect should fail
 
124
        client = self.get_client()
 
125
        self.assertRaises(socket.error,
 
126
                          client.connect, (server.host, server.port))
 
127
 
 
128
    def test_client_talks_server_respond(self):
 
129
        server = self.get_server()
 
130
        client = self.get_client()
 
131
        client.connect((server.host, server.port))
 
132
        self.assertIs(None, client.write('ping\n'))
 
133
        resp = client.read()
 
134
        self.assertClientAddr(client, server, 0)
 
135
        self.assertEquals('pong\n', resp)
 
136
 
 
137
    def test_server_fails_to_start(self):
 
138
        class CantStart(Exception):
 
139
            pass
 
140
 
 
141
        class CantStartServer(test_server.TestingTCPServer):
 
142
 
 
143
            def server_bind(self):
 
144
                raise CantStart()
 
145
 
 
146
        # The exception is raised in the main thread
 
147
        self.assertRaises(CantStart,
 
148
                          self.get_server, server_class=CantStartServer)
 
149
 
 
150
    def test_server_fails_while_serving_or_stoping(self):
 
151
        class ServerFailure(Exception):
 
152
            pass
 
153
 
 
154
        class FailingConnectionHandler(TCPConnectionHandler):
 
155
 
 
156
            def handle(self):
 
157
                raise ServerFailure()
 
158
 
 
159
        server = self.get_server(
 
160
            connection_handler_class=FailingConnectionHandler)
 
161
        # The server won't fail until a client connect
 
162
        client = self.get_client()
 
163
        client.connect((server.host, server.port))
 
164
        try:
 
165
            # Now we must force the server to answer by sending the request and
 
166
            # waiting for some answer. But since we don't control when the
 
167
            # server thread will be given cycles, we don't control either
 
168
            # whether our reads or writes may hang.
 
169
            client.sock.settimeout(0.1)
 
170
            client.write('ping\n')
 
171
            client.read()
 
172
        except socket.error:
 
173
            pass
 
174
        # Now the server has raise the exception in its own thread
 
175
        self.assertRaises(ServerFailure, server.stop_server)
 
176
 
 
177