Lib/test/test_socket.py

   1 #!/usr/bin/env python
   2
   3 import unittest
   4 import test_support
   5
   6 import socket
   7 import select
   8 import time
   9 import thread, threading
  10 import Queue
  11
  12 PORT = 50007
  13 HOST = 'localhost'
  14 MSG = 'Michael Gilfix was here\n'
  15
  16 class SocketTCPTest(unittest.TestCase):
  17
  18     def setUp(self):
  19         self.serv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  20         self.serv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  21         self.serv.bind((HOST, PORT))
  22         self.serv.listen(1)
  23
  24     def tearDown(self):
  25         self.serv.close()
  26         self.serv = None
  27
  28 class SocketUDPTest(unittest.TestCase):
  29
  30     def setUp(self):
  31         self.serv = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  32         self.serv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  33         self.serv.bind((HOST, PORT))
  34
  35     def tearDown(self):
  36         self.serv.close()
  37         self.serv = None
  38
  39 class ThreadableTest:
  40     """Threadable Test class
  41
  42     The ThreadableTest class makes it easy to create a threaded
  43     client/server pair from an existing unit test. To create a
  44     new threaded class from an existing unit test, use multiple
  45     inheritance:
  46
  47         class NewClass (OldClass, ThreadableTest):
  48             pass
  49
  50     This class defines two new fixture functions with obvious
  51     purposes for overriding:
  52
  53         clientSetUp ()
  54         clientTearDown ()
  55
  56     Any new test functions within the class must then define
  57     tests in pairs, where the test name is preceeded with a
  58     '_' to indicate the client portion of the test. Ex:
  59
  60         def testFoo(self):
  61             # Server portion
  62
  63         def _testFoo(self):
  64             # Client portion
  65
  66     Any exceptions raised by the clients during their tests
  67     are caught and transferred to the main thread to alert
  68     the testing framework.
  69
  70     Note, the server setup function cannot call any blocking
  71     functions that rely on the client thread during setup,
  72     unless serverExplicityReady() is called just before
  73     the blocking call (such as in setting up a client/server
  74     connection and performing the accept() in setUp().
  75     """
  76
  77     def __init__(self):
  78         # Swap the true setup function
  79         self.__setUp = self.setUp
  80         self.__tearDown = self.tearDown
  81         self.setUp = self._setUp
  82         self.tearDown = self._tearDown
  83
  84     def serverExplicitReady(self):
  85         """This method allows the server to explicitly indicate that
  86         it wants the client thread to proceed. This is useful if the
  87         server is about to execute a blocking routine that is
  88         dependent upon the client thread during its setup routine."""
  89         self.server_ready.set()
  90
  91     def _setUp(self):
  92         self.server_ready = threading.Event()
  93         self.client_ready = threading.Event()
  94         self.done = threading.Event()
  95         self.queue = Queue.Queue(1)
  96
  97         # Do some munging to start the client test.
  98         methodname = self.id()
  99         i = methodname.rfind('.')
 100         methodname = methodname[i+1:]
 101         test_method = getattr(self, '_' + methodname)
 102         self.client_thread = thread.start_new_thread(
 103             self.clientRun, (test_method,))
 104
 105         self.__setUp()
 106         if not self.server_ready.isSet():
 107             self.server_ready.set()
 108         self.client_ready.wait()
 109
 110     def _tearDown(self):
 111         self.__tearDown()
 112         self.done.wait()
 113
 114         if not self.queue.empty():
 115             msg = self.queue.get()
 116             self.fail(msg)
 117
 118     def clientRun(self, test_func):
 119         self.server_ready.wait()
 120         self.client_ready.set()
 121         self.clientSetUp()
 122         if not callable(test_func):
 123             raise TypeError, "test_func must be a callable function"
 124         try:
 125             test_func()
 126         except Exception, strerror:
 127             self.queue.put(strerror)
 128         self.clientTearDown()
 129
 130     def clientSetUp(self):
 131         raise NotImplementedError, "clientSetUp must be implemented."
 132
 133     def clientTearDown(self):
 134         self.done.set()
 135         thread.exit()
 136
 137 class ThreadedTCPSocketTest(SocketTCPTest, ThreadableTest):
 138
 139     def __init__(self, methodName='runTest'):
 140         SocketTCPTest.__init__(self, methodName=methodName)
 141         ThreadableTest.__init__(self)
 142
 143     def clientSetUp(self):
 144         self.cli = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 145
 146     def clientTearDown(self):
 147         self.cli.close()
 148         self.cli = None
 149         ThreadableTest.clientTearDown(self)
 150
 151 class ThreadedUDPSocketTest(SocketUDPTest, ThreadableTest):
 152
 153     def __init__(self, methodName='runTest'):
 154         SocketUDPTest.__init__(self, methodName=methodName)
 155         ThreadableTest.__init__(self)
 156
 157     def clientSetUp(self):
 158         self.cli = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 159
 160 class SocketConnectedTest(ThreadedTCPSocketTest):
 161
 162     def __init__(self, methodName='runTest'):
 163         ThreadedTCPSocketTest.__init__(self, methodName=methodName)
 164
 165     def setUp(self):
 166         ThreadedTCPSocketTest.setUp(self)
 167         # Indicate explicitly we're ready for the client thread to
 168         # proceed and then perform the blocking call to accept
 169         self.serverExplicitReady()
 170         conn, addr = self.serv.accept()
 171         self.cli_conn = conn
 172
 173     def tearDown(self):
 174         self.cli_conn.close()
 175         self.cli_conn = None
 176         ThreadedTCPSocketTest.tearDown(self)
 177
 178     def clientSetUp(self):
 179         ThreadedTCPSocketTest.clientSetUp(self)
 180         self.cli.connect((HOST, PORT))
 181         self.serv_conn = self.cli
 182
 183     def clientTearDown(self):
 184         self.serv_conn.close()
 185         self.serv_conn = None
 186         ThreadedTCPSocketTest.clientTearDown(self)
 187
 188 #######################################################################
 189 ## Begin Tests
 190
 191 class GeneralModuleTests(unittest.TestCase):
 192
 193     def testSocketError(self):
 194         """Testing that socket module exceptions."""
 195         def raise_error(*args, **kwargs):
 196             raise socket.error
 197         def raise_herror(*args, **kwargs):
 198             raise socket.herror
 199         def raise_gaierror(*args, **kwargs):
 200             raise socket.gaierror
 201         self.failUnlessRaises(socket.error, raise_error,
 202                               "Error raising socket exception.")
 203         self.failUnlessRaises(socket.error, raise_herror,
 204                               "Error raising socket exception.")
 205         self.failUnlessRaises(socket.error, raise_gaierror,
 206                               "Error raising socket exception.")
 207
 208     def testCrucialConstants(self):
 209         """Testing for mission critical constants."""
 210         socket.AF_INET
 211         socket.SOCK_STREAM
 212         socket.SOCK_DGRAM
 213         socket.SOCK_RAW
 214         socket.SOCK_RDM
 215         socket.SOCK_SEQPACKET
 216         socket.SOL_SOCKET
 217         socket.SO_REUSEADDR
 218
 219     def testNonCrucialConstants(self):
 220         """Testing for existance of non-crucial constants."""
 221         for const in (
 222                  "AF_UNIX",
 223
 224                  "SO_DEBUG", "SO_ACCEPTCONN", "SO_REUSEADDR", "SO_KEEPALIVE",
 225                  "SO_DONTROUTE", "SO_BROADCAST", "SO_USELOOPBACK", "SO_LINGER",
 226                  "SO_OOBINLINE", "SO_REUSEPORT", "SO_SNDBUF", "SO_RCVBUF",
 227                  "SO_SNDLOWAT", "SO_RCVLOWAT", "SO_SNDTIMEO", "SO_RCVTIMEO",
 228                  "SO_ERROR", "SO_TYPE", "SOMAXCONN",
 229
 230                  "MSG_OOB", "MSG_PEEK", "MSG_DONTROUTE", "MSG_EOR",
 231                  "MSG_TRUNC", "MSG_CTRUNC", "MSG_WAITALL", "MSG_BTAG",
 232                  "MSG_ETAG",
 233
 234                  "SOL_SOCKET",
 235
 236                  "IPPROTO_IP", "IPPROTO_ICMP", "IPPROTO_IGMP",
 237                  "IPPROTO_GGP", "IPPROTO_TCP", "IPPROTO_EGP",
 238                  "IPPROTO_PUP", "IPPROTO_UDP", "IPPROTO_IDP",
 239                  "IPPROTO_HELLO", "IPPROTO_ND", "IPPROTO_TP",
 240                  "IPPROTO_XTP", "IPPROTO_EON", "IPPROTO_BIP",
 241                  "IPPROTO_RAW", "IPPROTO_MAX",
 242
 243                  "IPPORT_RESERVED", "IPPORT_USERRESERVED",
 244
 245                  "INADDR_ANY", "INADDR_BROADCAST", "INADDR_LOOPBACK",
 246                  "INADDR_UNSPEC_GROUP", "INADDR_ALLHOSTS_GROUP",
 247                  "INADDR_MAX_LOCAL_GROUP", "INADDR_NONE",
 248
 249                  "IP_OPTIONS", "IP_HDRINCL", "IP_TOS", "IP_TTL",
 250                  "IP_RECVOPTS", "IP_RECVRETOPTS", "IP_RECVDSTADDR",
 251                  "IP_RETOPTS", "IP_MULTICAST_IF", "IP_MULTICAST_TTL",
 252                  "IP_MULTICAST_LOOP", "IP_ADD_MEMBERSHIP",
 253                  "IP_DROP_MEMBERSHIP",
 254                   ):
 255             try:
 256                 getattr(socket, const)
 257             except AttributeError:
 258                 pass
 259
 260     def testHostnameRes(self):
 261         """Testing hostname resolution mechanisms."""
 262         hostname = socket.gethostname()
 263         ip = socket.gethostbyname(hostname)
 264         self.assert_(ip.find('.') >= 0, "Error resolving host to ip.")
 265         hname, aliases, ipaddrs = socket.gethostbyaddr(ip)
 266         all_host_names = [hname] + aliases
 267         fqhn = socket.getfqdn()
 268         if not fqhn in all_host_names:
 269             self.fail("Error testing host resolution mechanisms.")
 270
 271     def testRefCountGetNameInfo(self):
 272         """Testing reference count for getnameinfo."""
 273         import sys
 274         if hasattr(sys, "getrefcount"):
 275             try:
 276                 # On some versions, this loses a reference
 277                 orig = sys.getrefcount(__name__)
 278                 socket.getnameinfo(__name__,0)
 279             except SystemError:
 280                 if sys.getrefcount(__name__) <> orig:
 281                     self.fail("socket.getnameinfo loses a reference")
 282
 283     def testInterpreterCrash(self):
 284         """Making sure getnameinfo doesn't crash the interpreter."""
 285         try:
 286             # On some versions, this crashes the interpreter.
 287             socket.getnameinfo(('x', 0, 0, 0), 0)
 288         except socket.error:
 289             pass
 290
 291     def testGetServByName(self):
 292         """Testing getservbyname()."""
 293         if hasattr(socket, 'getservbyname'):
 294             socket.getservbyname('telnet', 'tcp')
 295             try:
 296                 socket.getservbyname('telnet', 'udp')
 297             except socket.error:
 298                 pass
 299
 300     def testSockName(self):
 301         """Testing getsockname()."""
 302         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 303         sock.bind(("0.0.0.0", PORT+1))
 304         name = sock.getsockname()
 305         self.assertEqual(name, ("0.0.0.0", PORT+1))
 306
 307     def testGetSockOpt(self):
 308         """Testing getsockopt()."""
 309         # We know a socket should start without reuse==0
 310         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 311         reuse = sock.getsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR)
 312         self.failIf(reuse != 0, "initial mode is reuse")
 313
 314     def testSetSockOpt(self):
 315         """Testing setsockopt()."""
 316         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 317         sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 318         reuse = sock.getsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR)
 319         self.failIf(reuse == 0, "failed to set reuse mode")
 320
 321 class BasicTCPTest(SocketConnectedTest):
 322
 323     def __init__(self, methodName='runTest'):
 324         SocketConnectedTest.__init__(self, methodName=methodName)
 325
 326     def testRecv(self):
 327         """Testing large receive over TCP."""
 328         msg = self.cli_conn.recv(1024)
 329         self.assertEqual(msg, MSG)
 330
 331     def _testRecv(self):
 332         self.serv_conn.send(MSG)
 333
 334     def testOverFlowRecv(self):
 335         """Testing receive in chunks over TCP."""
 336         seg1 = self.cli_conn.recv(len(MSG) - 3)
 337         seg2 = self.cli_conn.recv(1024)
 338         msg = seg1 + seg2
 339         self.assertEqual(msg, MSG)
 340
 341     def _testOverFlowRecv(self):
 342         self.serv_conn.send(MSG)
 343
 344     def testRecvFrom(self):
 345         """Testing large recvfrom() over TCP."""
 346         msg, addr = self.cli_conn.recvfrom(1024)
 347         hostname, port = addr
 348         ##self.assertEqual(hostname, socket.gethostbyname('localhost'))
 349         self.assertEqual(msg, MSG)
 350
 351     def _testRecvFrom(self):
 352         self.serv_conn.send(MSG)
 353
 354     def testOverFlowRecvFrom(self):
 355         """Testing recvfrom() in chunks over TCP."""
 356         seg1, addr = self.cli_conn.recvfrom(len(MSG)-3)
 357         seg2, addr = self.cli_conn.recvfrom(1024)
 358         msg = seg1 + seg2
 359         hostname, port = addr
 360         ##self.assertEqual(hostname, socket.gethostbyname('localhost'))
 361         self.assertEqual(msg, MSG)
 362
 363     def _testOverFlowRecvFrom(self):
 364         self.serv_conn.send(MSG)
 365
 366     def testSendAll(self):
 367         """Testing sendall() with a 2048 byte string over TCP."""
 368         while 1:
 369             read = self.cli_conn.recv(1024)
 370             if not read:
 371                 break
 372             self.assert_(len(read) == 1024, "Error performing sendall.")
 373             read = filter(lambda x: x == 'f', read)
 374             self.assert_(len(read) == 1024, "Error performing sendall.")
 375
 376     def _testSendAll(self):
 377         big_chunk = 'f' * 2048
 378         self.serv_conn.sendall(big_chunk)
 379
 380     def testFromFd(self):
 381         """Testing fromfd()."""
 382         if not hasattr(socket, "fromfd"):
 383             return # On Windows, this doesn't exist
 384         fd = self.cli_conn.fileno()
 385         sock = socket.fromfd(fd, socket.AF_INET, socket.SOCK_STREAM)
 386         msg = sock.recv(1024)
 387         self.assertEqual(msg, MSG)
 388
 389     def _testFromFd(self):
 390         self.serv_conn.send(MSG)
 391
 392     def testShutdown(self):
 393         """Testing shutdown()."""
 394         msg = self.cli_conn.recv(1024)
 395         self.assertEqual(msg, MSG)
 396
 397     def _testShutdown(self):
 398         self.serv_conn.send(MSG)
 399         self.serv_conn.shutdown(2)
 400
 401 class BasicUDPTest(ThreadedUDPSocketTest):
 402
 403     def __init__(self, methodName='runTest'):
 404         ThreadedUDPSocketTest.__init__(self, methodName=methodName)
 405
 406     def testSendtoAndRecv(self):
 407         """Testing sendto() and Recv() over UDP."""
 408         msg = self.serv.recv(len(MSG))
 409         self.assertEqual(msg, MSG)
 410
 411     def _testSendtoAndRecv(self):
 412         self.cli.sendto(MSG, 0, (HOST, PORT))
 413
 414     def testRecvFrom(self):
 415         """Testing recvfrom() over UDP."""
 416         msg, addr = self.serv.recvfrom(len(MSG))
 417         hostname, port = addr
 418         ##self.assertEqual(hostname, socket.gethostbyname('localhost'))
 419         self.assertEqual(msg, MSG)
 420
 421     def _testRecvFrom(self):
 422         self.cli.sendto(MSG, 0, (HOST, PORT))
 423
 424 class NonBlockingTCPTests(ThreadedTCPSocketTest):
 425
 426     def __init__(self, methodName='runTest'):
 427         ThreadedTCPSocketTest.__init__(self, methodName=methodName)
 428
 429     def testSetBlocking(self):
 430         """Testing whether set blocking works."""
 431         self.serv.setblocking(0)
 432         start = time.time()
 433         try:
 434             self.serv.accept()
 435         except socket.error:
 436             pass
 437         end = time.time()
 438         self.assert_((end - start) < 1.0, "Error setting non-blocking mode.")
 439
 440     def _testSetBlocking(self):
 441         pass
 442
 443     def testAccept(self):
 444         """Testing non-blocking accept."""
 445         self.serv.setblocking(0)
 446         try:
 447             conn, addr = self.serv.accept()
 448         except socket.error:
 449             pass
 450         else:
 451             self.fail("Error trying to do non-blocking accept.")
 452         read, write, err = select.select([self.serv], [], [])
 453         if self.serv in read:
 454             conn, addr = self.serv.accept()
 455         else:
 456             self.fail("Error trying to do accept after select.")
 457
 458     def _testAccept(self):
 459         self.cli.connect((HOST, PORT))
 460
 461     def testConnect(self):
 462         """Testing non-blocking connect."""
 463         conn, addr = self.serv.accept()
 464
 465     def _testConnect(self):
 466         self.cli.settimeout(10)
 467         self.cli.connect((HOST, PORT))
 468
 469     def testRecv(self):
 470         """Testing non-blocking recv."""
 471         conn, addr = self.serv.accept()
 472         conn.setblocking(0)
 473         try:
 474             msg = conn.recv(len(MSG))
 475         except socket.error:
 476             pass
 477         else:
 478             self.fail("Error trying to do non-blocking recv.")
 479         read, write, err = select.select([conn], [], [])
 480         if conn in read:
 481             msg = conn.recv(len(MSG))
 482             self.assertEqual(msg, MSG)
 483         else:
 484             self.fail("Error during select call to non-blocking socket.")
 485
 486     def _testRecv(self):
 487         self.cli.connect((HOST, PORT))
 488         self.cli.send(MSG)
 489
 490 class FileObjectClassTestCase(SocketConnectedTest):
 491
 492     def __init__(self, methodName='runTest'):
 493         SocketConnectedTest.__init__(self, methodName=methodName)
 494
 495     def setUp(self):
 496         SocketConnectedTest.setUp(self)
 497         self.serv_file = socket._fileobject(self.cli_conn, 'rb', 8192)
 498
 499     def tearDown(self):
 500         self.serv_file.close()
 501         self.serv_file = None
 502         SocketConnectedTest.tearDown(self)
 503
 504     def clientSetUp(self):
 505         SocketConnectedTest.clientSetUp(self)
 506         self.cli_file = socket._fileobject(self.serv_conn, 'rb', 8192)
 507
 508     def clientTearDown(self):
 509         self.cli_file.close()
 510         self.cli_file = None
 511         SocketConnectedTest.clientTearDown(self)
 512
 513     def testSmallRead(self):
 514         """Performing small file read test."""
 515         first_seg = self.serv_file.read(len(MSG)-3)
 516         second_seg = self.serv_file.read(3)
 517         msg = first_seg + second_seg
 518         self.assertEqual(msg, MSG)
 519
 520     def _testSmallRead(self):
 521         self.cli_file.write(MSG)
 522         self.cli_file.flush()
 523
 524     def testUnbufferedRead(self):
 525         """Performing unbuffered file read test."""
 526         buf = ''
 527         while 1:
 528             char = self.serv_file.read(1)
 529             self.failIf(not char)
 530             buf += char
 531             if buf == MSG:
 532                 break
 533
 534     def _testUnbufferedRead(self):
 535         self.cli_file.write(MSG)
 536         self.cli_file.flush()
 537
 538     def testReadline(self):
 539         """Performing file readline test."""
 540         line = self.serv_file.readline()
 541         self.assertEqual(line, MSG)
 542
 543     def _testReadline(self):
 544         self.cli_file.write(MSG)
 545         self.cli_file.flush()
 546
 547 def main():
 548     suite = unittest.TestSuite()
 549     suite.addTest(unittest.makeSuite(GeneralModuleTests))
 550     suite.addTest(unittest.makeSuite(BasicTCPTest))
 551     suite.addTest(unittest.makeSuite(BasicUDPTest))
 552     suite.addTest(unittest.makeSuite(NonBlockingTCPTests))
 553     suite.addTest(unittest.makeSuite(FileObjectClassTestCase))
 554     test_support.run_suite(suite)
 555
 556 if __name__ == "__main__":
 557     main()