twisted/test/proto_helpers.py

   1 # -*- test-case-name: twisted.test.test_stringtransport -*-
   2 # Copyright (c) Twisted Matrix Laboratories.
   3 # See LICENSE for details.
   4
   5 """
   6 Assorted functionality which is commonly useful when writing unit tests.
   7 """
   8
   9 from socket import AF_INET, AF_INET6
  10 from StringIO import StringIO
  11
  12 from zope.interface import implements
  13
  14 from twisted.internet.interfaces import (
  15     ITransport, IConsumer, IPushProducer, IConnector)
  16 from twisted.internet.interfaces import (
  17     IReactorTCP, IReactorSSL, IReactorUNIX, IReactorSocket)
  18 from twisted.internet.interfaces import IListeningPort
  19 from twisted.protocols import basic
  20 from twisted.internet import protocol, error, address
  21
  22 from twisted.internet.address import IPv4Address, UNIXAddress
  23
  24
  25 class AccumulatingProtocol(protocol.Protocol):
  26     """
  27     L{AccumulatingProtocol} is an L{IProtocol} implementation which collects
  28     the data delivered to it and can fire a Deferred when it is connected or
  29     disconnected.
  30
  31     @ivar made: A flag indicating whether C{connectionMade} has been called.
  32     @ivar data: A string giving all the data passed to C{dataReceived}.
  33     @ivar closed: A flag indicated whether C{connectionLost} has been called.
  34     @ivar closedReason: The value of the I{reason} parameter passed to
  35         C{connectionLost}.
  36     @ivar closedDeferred: If set to a L{Deferred}, this will be fired when
  37         C{connectionLost} is called.
  38     """
  39     made = closed = 0
  40     closedReason = None
  41
  42     closedDeferred = None
  43
  44     data = ""
  45
  46     factory = None
  47
  48     def connectionMade(self):
  49         self.made = 1
  50         if (self.factory is not None and
  51             self.factory.protocolConnectionMade is not None):
  52             d = self.factory.protocolConnectionMade
  53             self.factory.protocolConnectionMade = None
  54             d.callback(self)
  55
  56     def dataReceived(self, data):
  57         self.data += data
  58
  59     def connectionLost(self, reason):
  60         self.closed = 1
  61         self.closedReason = reason
  62         if self.closedDeferred is not None:
  63             d, self.closedDeferred = self.closedDeferred, None
  64             d.callback(None)
  65
  66
  67 class LineSendingProtocol(basic.LineReceiver):
  68     lostConn = False
  69
  70     def __init__(self, lines, start = True):
  71         self.lines = lines[:]
  72         self.response = []
  73         self.start = start
  74
  75     def connectionMade(self):
  76         if self.start:
  77             map(self.sendLine, self.lines)
  78
  79     def lineReceived(self, line):
  80         if not self.start:
  81             map(self.sendLine, self.lines)
  82             self.lines = []
  83         self.response.append(line)
  84
  85     def connectionLost(self, reason):
  86         self.lostConn = True
  87
  88
  89 class FakeDatagramTransport:
  90     noAddr = object()
  91
  92     def __init__(self):
  93         self.written = []
  94
  95     def write(self, packet, addr=noAddr):
  96         self.written.append((packet, addr))
  97
  98
  99 class StringTransport:
 100     """
 101     A transport implementation which buffers data in memory and keeps track of
 102     its other state without providing any behavior.
 103
 104     L{StringTransport} has a number of attributes which are not part of any of
 105     the interfaces it claims to implement.  These attributes are provided for
 106     testing purposes.  Implementation code should not use any of these
 107     attributes; they are not provided by other transports.
 108
 109     @ivar disconnecting: A C{bool} which is C{False} until L{loseConnection} is
 110         called, then C{True}.
 111
 112     @ivar producer: If a producer is currently registered, C{producer} is a
 113         reference to it.  Otherwise, C{None}.
 114
 115     @ivar streaming: If a producer is currently registered, C{streaming} refers
 116         to the value of the second parameter passed to C{registerProducer}.
 117
 118     @ivar hostAddr: C{None} or an object which will be returned as the host
 119         address of this transport.  If C{None}, a nasty tuple will be returned
 120         instead.
 121
 122     @ivar peerAddr: C{None} or an object which will be returned as the peer
 123         address of this transport.  If C{None}, a nasty tuple will be returned
 124         instead.
 125
 126     @ivar producerState: The state of this L{StringTransport} in its capacity
 127         as an L{IPushProducer}.  One of C{'producing'}, C{'paused'}, or
 128         C{'stopped'}.
 129
 130     @ivar io: A L{StringIO} which holds the data which has been written to this
 131         transport since the last call to L{clear}.  Use L{value} instead of
 132         accessing this directly.
 133     """
 134     implements(ITransport, IConsumer, IPushProducer)
 135
 136     disconnecting = False
 137
 138     producer = None
 139     streaming = None
 140
 141     hostAddr = None
 142     peerAddr = None
 143
 144     producerState = 'producing'
 145
 146     def __init__(self, hostAddress=None, peerAddress=None):
 147         self.clear()
 148         if hostAddress is not None:
 149             self.hostAddr = hostAddress
 150         if peerAddress is not None:
 151             self.peerAddr = peerAddress
 152         self.connected = True
 153
 154     def clear(self):
 155         """
 156         Discard all data written to this transport so far.
 157
 158         This is not a transport method.  It is intended for tests.  Do not use
 159         it in implementation code.
 160         """
 161         self.io = StringIO()
 162
 163
 164     def value(self):
 165         """
 166         Retrieve all data which has been buffered by this transport.
 167
 168         This is not a transport method.  It is intended for tests.  Do not use
 169         it in implementation code.
 170
 171         @return: A C{str} giving all data written to this transport since the
 172             last call to L{clear}.
 173         @rtype: C{str}
 174         """
 175         return self.io.getvalue()
 176
 177
 178     # ITransport
 179     def write(self, data):
 180         if isinstance(data, unicode): # no, really, I mean it
 181             raise TypeError("Data must not be unicode")
 182         self.io.write(data)
 183
 184
 185     def writeSequence(self, data):
 186         self.io.write(''.join(data))
 187
 188
 189     def loseConnection(self):
 190         """
 191         Close the connection. Does nothing besides toggle the C{disconnecting}
 192         instance variable to C{True}.
 193         """
 194         self.disconnecting = True
 195
 196
 197     def getPeer(self):
 198         if self.peerAddr is None:
 199             return address.IPv4Address('TCP', '192.168.1.1', 54321)
 200         return self.peerAddr
 201
 202
 203     def getHost(self):
 204         if self.hostAddr is None:
 205             return address.IPv4Address('TCP', '10.0.0.1', 12345)
 206         return self.hostAddr
 207
 208
 209     # IConsumer
 210     def registerProducer(self, producer, streaming):
 211         if self.producer is not None:
 212             raise RuntimeError("Cannot register two producers")
 213         self.producer = producer
 214         self.streaming = streaming
 215
 216
 217     def unregisterProducer(self):
 218         if self.producer is None:
 219             raise RuntimeError(
 220                 "Cannot unregister a producer unless one is registered")
 221         self.producer = None
 222         self.streaming = None
 223
 224
 225     # IPushProducer
 226     def _checkState(self):
 227         if self.disconnecting:
 228             raise RuntimeError(
 229                 "Cannot resume producing after loseConnection")
 230         if self.producerState == 'stopped':
 231             raise RuntimeError("Cannot resume a stopped producer")
 232
 233
 234     def pauseProducing(self):
 235         self._checkState()
 236         self.producerState = 'paused'
 237
 238
 239     def stopProducing(self):
 240         self.producerState = 'stopped'
 241
 242
 243     def resumeProducing(self):
 244         self._checkState()
 245         self.producerState = 'producing'
 246
 247
 248
 249 class StringTransportWithDisconnection(StringTransport):
 250     def loseConnection(self):
 251         if self.connected:
 252             self.connected = False
 253             self.protocol.connectionLost(error.ConnectionDone("Bye."))
 254
 255
 256
 257 class StringIOWithoutClosing(StringIO):
 258     """
 259     A StringIO that can't be closed.
 260     """
 261     def close(self):
 262         """
 263         Do nothing.
 264         """
 265
 266
 267
 268 class _FakePort(object):
 269     """
 270     A fake L{IListeningPort} to be used in tests.
 271
 272     @ivar _hostAddress: The L{IAddress} this L{IListeningPort} is pretending
 273         to be listening on.
 274     """
 275     implements(IListeningPort)
 276
 277     def __init__(self, hostAddress):
 278         """
 279         @param hostAddress: An L{IAddress} this L{IListeningPort} should
 280             pretend to be listening on.
 281         """
 282         self._hostAddress = hostAddress
 283
 284
 285     def startListening(self):
 286         """
 287         Fake L{IListeningPort.startListening} that doesn't do anything.
 288         """
 289
 290
 291     def stopListening(self):
 292         """
 293         Fake L{IListeningPort.stopListening} that doesn't do anything.
 294         """
 295
 296
 297     def getHost(self):
 298         """
 299         Fake L{IListeningPort.getHost} that returns our L{IAddress}.
 300         """
 301         return self._hostAddress
 302
 303
 304
 305 class _FakeConnector(object):
 306     """
 307     A fake L{IConnector} that allows us to inspect if it has been told to stop
 308     connecting.
 309
 310     @ivar stoppedConnecting: has this connector's
 311         L{FakeConnector.stopConnecting} method been invoked yet?
 312
 313     @ivar _address: An L{IAddress} provider that represents our destination.
 314     """
 315     implements(IConnector)
 316
 317     stoppedConnecting = False
 318
 319     def __init__(self, address):
 320         """
 321         @param address: An L{IAddress} provider that represents this
 322             connector's destination.
 323         """
 324         self._address = address
 325
 326
 327     def stopConnecting(self):
 328         """
 329         Implement L{IConnector.stopConnecting} and set
 330         L{FakeConnector.stoppedConnecting} to C{True}
 331         """
 332         self.stoppedConnecting = True
 333
 334
 335     def disconnect(self):
 336         """
 337         Implement L{IConnector.disconnect} as a no-op.
 338         """
 339
 340
 341     def connect(self):
 342         """
 343         Implement L{IConnector.connect} as a no-op.
 344         """
 345
 346
 347     def getDestination(self):
 348         """
 349         Implement L{IConnector.getDestination} to return the C{address} passed
 350         to C{__init__}.
 351         """
 352         return self._address
 353
 354
 355
 356 class MemoryReactor(object):
 357     """
 358     A fake reactor to be used in tests.  This reactor doesn't actually do
 359     much that's useful yet.  It accepts TCP connection setup attempts, but
 360     they will never succeed.
 361
 362     @ivar tcpClients: a list that keeps track of connection attempts (ie, calls
 363         to C{connectTCP}).
 364     @type tcpClients: C{list}
 365
 366     @ivar tcpServers: a list that keeps track of server listen attempts (ie, calls
 367         to C{listenTCP}).
 368     @type tcpServers: C{list}
 369
 370     @ivar sslClients: a list that keeps track of connection attempts (ie,
 371         calls to C{connectSSL}).
 372     @type sslClients: C{list}
 373
 374     @ivar sslServers: a list that keeps track of server listen attempts (ie,
 375         calls to C{listenSSL}).
 376     @type sslServers: C{list}
 377
 378     @ivar unixClients: a list that keeps track of connection attempts (ie,
 379         calls to C{connectUNIX}).
 380     @type unixClients: C{list}
 381
 382     @ivar unixServers: a list that keeps track of server listen attempts (ie,
 383         calls to C{listenUNIX}).
 384     @type unixServers: C{list}
 385
 386     @ivar adoptedPorts: a list that keeps track of server listen attempts (ie,
 387         calls to C{adoptStreamPort}).
 388     """
 389     implements(IReactorTCP, IReactorSSL, IReactorUNIX, IReactorSocket)
 390
 391     def __init__(self):
 392         """
 393         Initialize the tracking lists.
 394         """
 395         self.tcpClients = []
 396         self.tcpServers = []
 397         self.sslClients = []
 398         self.sslServers = []
 399         self.unixClients = []
 400         self.unixServers = []
 401         self.adoptedPorts = []
 402
 403
 404     def adoptStreamPort(self, fileno, addressFamily, factory):
 405         """
 406         Fake L{IReactorSocket.adoptStreamPort}, that logs the call and returns
 407         an L{IListeningPort}.
 408         """
 409         if addressFamily == AF_INET:
 410             addr = IPv4Address('TCP', '0.0.0.0', 1234)
 411         elif addressFamily == AF_INET6:
 412             addr = IPv6Address('TCP', '::', 1234)
 413         else:
 414             raise UnsupportedAddressFamily()
 415
 416         self.adoptedPorts.append((fileno, addressFamily, factory))
 417         return _FakePort(addr)
 418
 419
 420     def listenTCP(self, port, factory, backlog=50, interface=''):
 421         """
 422         Fake L{reactor.listenTCP}, that logs the call and returns an
 423         L{IListeningPort}.
 424         """
 425         self.tcpServers.append((port, factory, backlog, interface))
 426         return _FakePort(IPv4Address('TCP', '0.0.0.0', port))
 427
 428
 429     def connectTCP(self, host, port, factory, timeout=30, bindAddress=None):
 430         """
 431         Fake L{reactor.connectTCP}, that logs the call and returns an
 432         L{IConnector}.
 433         """
 434         self.tcpClients.append((host, port, factory, timeout, bindAddress))
 435         conn = _FakeConnector(IPv4Address('TCP', host, port))
 436         factory.startedConnecting(conn)
 437         return conn
 438
 439
 440     def listenSSL(self, port, factory, contextFactory,
 441                   backlog=50, interface=''):
 442         """
 443         Fake L{reactor.listenSSL}, that logs the call and returns an
 444         L{IListeningPort}.
 445         """
 446         self.sslServers.append((port, factory, contextFactory,
 447                                 backlog, interface))
 448         return _FakePort(IPv4Address('TCP', '0.0.0.0', port))
 449
 450
 451     def connectSSL(self, host, port, factory, contextFactory,
 452                    timeout=30, bindAddress=None):
 453         """
 454         Fake L{reactor.connectSSL}, that logs the call and returns an
 455         L{IConnector}.
 456         """
 457         self.sslClients.append((host, port, factory, contextFactory,
 458                                 timeout, bindAddress))
 459         conn = _FakeConnector(IPv4Address('TCP', host, port))
 460         factory.startedConnecting(conn)
 461         return conn
 462
 463
 464     def listenUNIX(self, address, factory,
 465                    backlog=50, mode=0666, wantPID=0):
 466         """
 467         Fake L{reactor.listenUNIX}, that logs the call and returns an
 468         L{IListeningPort}.
 469         """
 470         self.unixServers.append((address, factory, backlog, mode, wantPID))
 471         return _FakePort(UNIXAddress(address))
 472
 473
 474     def connectUNIX(self, address, factory, timeout=30, checkPID=0):
 475         """
 476         Fake L{reactor.connectUNIX}, that logs the call and returns an
 477         L{IConnector}.
 478         """
 479         self.unixClients.append((address, factory, timeout, checkPID))
 480         conn = _FakeConnector(UNIXAddress(address))
 481         factory.startedConnecting(conn)
 482         return conn
 483
 484
 485
 486 class RaisingMemoryReactor(object):
 487     """
 488     A fake reactor to be used in tests.  It accepts TCP connection setup
 489     attempts, but they will fail.
 490
 491     @ivar _listenException: An instance of an L{Exception}
 492     @ivar _connectException: An instance of an L{Exception}
 493     """
 494     implements(IReactorTCP, IReactorSSL, IReactorUNIX, IReactorSocket)
 495
 496     def __init__(self, listenException=None, connectException=None):
 497         """
 498         @param listenException: An instance of an L{Exception} to raise when any
 499             C{listen} method is called.
 500
 501         @param connectException: An instance of an L{Exception} to raise when
 502             any C{connect} method is called.
 503         """
 504         self._listenException = listenException
 505         self._connectException = connectException
 506
 507
 508     def adoptStreamPort(self, fileno, addressFamily, factory):
 509         """
 510         Fake L{IReactorSocket.adoptStreamPort}, that raises
 511         L{self._listenException}.
 512         """
 513         raise self._listenException
 514
 515
 516     def listenTCP(self, port, factory, backlog=50, interface=''):
 517         """
 518         Fake L{reactor.listenTCP}, that raises L{self._listenException}.
 519         """
 520         raise self._listenException
 521
 522
 523     def connectTCP(self, host, port, factory, timeout=30, bindAddress=None):
 524         """
 525         Fake L{reactor.connectTCP}, that raises L{self._connectException}.
 526         """
 527         raise self._connectException
 528
 529
 530     def listenSSL(self, port, factory, contextFactory,
 531                   backlog=50, interface=''):
 532         """
 533         Fake L{reactor.listenSSL}, that raises L{self._listenException}.
 534         """
 535         raise self._listenException
 536
 537
 538     def connectSSL(self, host, port, factory, contextFactory,
 539                    timeout=30, bindAddress=None):
 540         """
 541         Fake L{reactor.connectSSL}, that raises L{self._connectException}.
 542         """
 543         raise self._connectException
 544
 545
 546     def listenUNIX(self, address, factory,
 547                    backlog=50, mode=0666, wantPID=0):
 548         """
 549         Fake L{reactor.listenUNIX}, that raises L{self._listenException}.
 550         """
 551         raise self._listenException
 552
 553
 554     def connectUNIX(self, address, factory, timeout=30, checkPID=0):
 555         """
 556         Fake L{reactor.connectUNIX}, that raises L{self._connectException}.
 557         """
 558         raise self._connectException