Lib/test/test_socketserver.py

   1 """
   2 Test suite for SocketServer.py.
   3 """
   4
   5 import contextlib
   6 import errno
   7 import imp
   8 import os
   9 import select
  10 import signal
  11 import socket
  12 import tempfile
  13 import threading
  14 import time
  15 import unittest
  16 import SocketServer
  17
  18 import test.test_support
  19 from test.test_support import reap_children, reap_threads, verbose
  20 from test.test_support import TESTFN as TEST_FILE
  21
  22 test.test_support.requires("network")
  23
  24 TEST_STR = "hello world\n"
  25 HOST = test.test_support.HOST
  26
  27 HAVE_UNIX_SOCKETS = hasattr(socket, "AF_UNIX")
  28 HAVE_FORKING = hasattr(os, "fork") and os.name != "os2"
  29
  30 def signal_alarm(n):
  31     """Call signal.alarm when it exists (i.e. not on Windows)."""
  32     if hasattr(signal, 'alarm'):
  33         signal.alarm(n)
  34
  35 def receive(sock, n, timeout=20):
  36     r, w, x = select.select([sock], [], [], timeout)
  37     if sock in r:
  38         return sock.recv(n)
  39     else:
  40         raise RuntimeError, "timed out on %r" % (sock,)
  41
  42 if HAVE_UNIX_SOCKETS:
  43     class ForkingUnixStreamServer(SocketServer.ForkingMixIn,
  44                                   SocketServer.UnixStreamServer):
  45         pass
  46
  47     class ForkingUnixDatagramServer(SocketServer.ForkingMixIn,
  48                                     SocketServer.UnixDatagramServer):
  49         pass
  50
  51
  52 @contextlib.contextmanager
  53 def simple_subprocess(testcase):
  54     pid = os.fork()
  55     if pid == 0:
  56         # Don't throw an exception; it would be caught by the test harness.
  57         os._exit(72)
  58     yield None
  59     pid2, status = os.waitpid(pid, 0)
  60     testcase.assertEquals(pid2, pid)
  61     testcase.assertEquals(72 << 8, status)
  62
  63
  64 class SocketServerTest(unittest.TestCase):
  65     """Test all socket servers."""
  66
  67     def setUp(self):
  68         signal_alarm(20)  # Kill deadlocks after 20 seconds.
  69         self.port_seed = 0
  70         self.test_files = []
  71
  72     def tearDown(self):
  73         signal_alarm(0)  # Didn't deadlock.
  74         reap_children()
  75
  76         for fn in self.test_files:
  77             try:
  78                 os.remove(fn)
  79             except os.error:
  80                 pass
  81         self.test_files[:] = []
  82
  83     def pickaddr(self, proto):
  84         if proto == socket.AF_INET:
  85             return (HOST, 0)
  86         else:
  87             # XXX: We need a way to tell AF_UNIX to pick its own name
  88             # like AF_INET provides port==0.
  89             dir = None
  90             if os.name == 'os2':
  91                 dir = '\socket'
  92             fn = tempfile.mktemp(prefix='unix_socket.', dir=dir)
  93             if os.name == 'os2':
  94                 # AF_UNIX socket names on OS/2 require a specific prefix
  95                 # which can't include a drive letter and must also use
  96                 # backslashes as directory separators
  97                 if fn[1] == ':':
  98                     fn = fn[2:]
  99                 if fn[0] in (os.sep, os.altsep):
 100                     fn = fn[1:]
 101                 if os.sep == '/':
 102                     fn = fn.replace(os.sep, os.altsep)
 103                 else:
 104                     fn = fn.replace(os.altsep, os.sep)
 105             self.test_files.append(fn)
 106             return fn
 107
 108     def make_server(self, addr, svrcls, hdlrbase):
 109         class MyServer(svrcls):
 110             def handle_error(self, request, client_address):
 111                 self.close_request(request)
 112                 self.server_close()
 113                 raise
 114
 115         class MyHandler(hdlrbase):
 116             def handle(self):
 117                 line = self.rfile.readline()
 118                 self.wfile.write(line)
 119
 120         if verbose: print "creating server"
 121         server = MyServer(addr, MyHandler)
 122         self.assertEquals(server.server_address, server.socket.getsockname())
 123         return server
 124
 125     @reap_threads
 126     def run_server(self, svrcls, hdlrbase, testfunc):
 127         server = self.make_server(self.pickaddr(svrcls.address_family),
 128                                   svrcls, hdlrbase)
 129         # We had the OS pick a port, so pull the real address out of
 130         # the server.
 131         addr = server.server_address
 132         if verbose:
 133             print "server created"
 134             print "ADDR =", addr
 135             print "CLASS =", svrcls
 136         t = threading.Thread(
 137             name='%s serving' % svrcls,
 138             target=server.serve_forever,
 139             # Short poll interval to make the test finish quickly.
 140             # Time between requests is short enough that we won't wake
 141             # up spuriously too many times.
 142             kwargs={'poll_interval':0.01})
 143         t.daemon = True  # In case this function raises.
 144         t.start()
 145         if verbose: print "server running"
 146         for i in range(3):
 147             if verbose: print "test client", i
 148             testfunc(svrcls.address_family, addr)
 149         if verbose: print "waiting for server"
 150         server.shutdown()
 151         t.join()
 152         if verbose: print "done"
 153
 154     def stream_examine(self, proto, addr):
 155         s = socket.socket(proto, socket.SOCK_STREAM)
 156         s.connect(addr)
 157         s.sendall(TEST_STR)
 158         buf = data = receive(s, 100)
 159         while data and '\n' not in buf:
 160             data = receive(s, 100)
 161             buf += data
 162         self.assertEquals(buf, TEST_STR)
 163         s.close()
 164
 165     def dgram_examine(self, proto, addr):
 166         s = socket.socket(proto, socket.SOCK_DGRAM)
 167         s.sendto(TEST_STR, addr)
 168         buf = data = receive(s, 100)
 169         while data and '\n' not in buf:
 170             data = receive(s, 100)
 171             buf += data
 172         self.assertEquals(buf, TEST_STR)
 173         s.close()
 174
 175     def test_TCPServer(self):
 176         self.run_server(SocketServer.TCPServer,
 177                         SocketServer.StreamRequestHandler,
 178                         self.stream_examine)
 179
 180     def test_ThreadingTCPServer(self):
 181         self.run_server(SocketServer.ThreadingTCPServer,
 182                         SocketServer.StreamRequestHandler,
 183                         self.stream_examine)
 184
 185     if HAVE_FORKING:
 186         def test_ForkingTCPServer(self):
 187             with simple_subprocess(self):
 188                 self.run_server(SocketServer.ForkingTCPServer,
 189                                 SocketServer.StreamRequestHandler,
 190                                 self.stream_examine)
 191
 192     if HAVE_UNIX_SOCKETS:
 193         def test_UnixStreamServer(self):
 194             self.run_server(SocketServer.UnixStreamServer,
 195                             SocketServer.StreamRequestHandler,
 196                             self.stream_examine)
 197
 198         def test_ThreadingUnixStreamServer(self):
 199             self.run_server(SocketServer.ThreadingUnixStreamServer,
 200                             SocketServer.StreamRequestHandler,
 201                             self.stream_examine)
 202
 203         if HAVE_FORKING:
 204             def test_ForkingUnixStreamServer(self):
 205                 with simple_subprocess(self):
 206                     self.run_server(ForkingUnixStreamServer,
 207                                     SocketServer.StreamRequestHandler,
 208                                     self.stream_examine)
 209
 210     def test_UDPServer(self):
 211         self.run_server(SocketServer.UDPServer,
 212                         SocketServer.DatagramRequestHandler,
 213                         self.dgram_examine)
 214
 215     def test_ThreadingUDPServer(self):
 216         self.run_server(SocketServer.ThreadingUDPServer,
 217                         SocketServer.DatagramRequestHandler,
 218                         self.dgram_examine)
 219
 220     if HAVE_FORKING:
 221         def test_ForkingUDPServer(self):
 222             with simple_subprocess(self):
 223                 self.run_server(SocketServer.ForkingUDPServer,
 224                                 SocketServer.DatagramRequestHandler,
 225                                 self.dgram_examine)
 226
 227     # Alas, on Linux (at least) recvfrom() doesn't return a meaningful
 228     # client address so this cannot work:
 229
 230     # if HAVE_UNIX_SOCKETS:
 231     #     def test_UnixDatagramServer(self):
 232     #         self.run_server(SocketServer.UnixDatagramServer,
 233     #                         SocketServer.DatagramRequestHandler,
 234     #                         self.dgram_examine)
 235     #
 236     #     def test_ThreadingUnixDatagramServer(self):
 237     #         self.run_server(SocketServer.ThreadingUnixDatagramServer,
 238     #                         SocketServer.DatagramRequestHandler,
 239     #                         self.dgram_examine)
 240     #
 241     #     if HAVE_FORKING:
 242     #         def test_ForkingUnixDatagramServer(self):
 243     #             self.run_server(SocketServer.ForkingUnixDatagramServer,
 244     #                             SocketServer.DatagramRequestHandler,
 245     #                             self.dgram_examine)
 246
 247
 248 def test_main():
 249     if imp.lock_held():
 250         # If the import lock is held, the threads will hang
 251         raise unittest.SkipTest("can't run when import lock is held")
 252
 253     test.test_support.run_unittest(SocketServerTest)
 254
 255 if __name__ == "__main__":
 256     test_main()
 257     signal_alarm(3)  # Shutdown shouldn't take more than 3 seconds.