twisted/test/test_tpfile.py

   1
   2 # Copyright (c) Twisted Matrix Laboratories.
   3 # See LICENSE for details.
   4
   5
   6 from twisted.trial import unittest
   7 from twisted.protocols import loopback
   8 from twisted.protocols import basic
   9 from twisted.internet import protocol, abstract
  10
  11 import StringIO
  12
  13 class BufferingServer(protocol.Protocol):
  14     buffer = ''
  15     def dataReceived(self, data):
  16         self.buffer += data
  17
  18 class FileSendingClient(protocol.Protocol):
  19     def __init__(self, f):
  20         self.f = f
  21
  22     def connectionMade(self):
  23         s = basic.FileSender()
  24         d = s.beginFileTransfer(self.f, self.transport, lambda x: x)
  25         d.addCallback(lambda r: self.transport.loseConnection())
  26
  27 class FileSenderTestCase(unittest.TestCase):
  28     def testSendingFile(self):
  29         testStr = 'xyz' * 100 + 'abc' * 100 + '123' * 100
  30         s = BufferingServer()
  31         c = FileSendingClient(StringIO.StringIO(testStr))
  32
  33         d = loopback.loopbackTCP(s, c)
  34         d.addCallback(lambda x : self.assertEqual(s.buffer, testStr))
  35         return d
  36
  37     def testSendingEmptyFile(self):
  38         fileSender = basic.FileSender()
  39         consumer = abstract.FileDescriptor()
  40         consumer.connected = 1
  41         emptyFile = StringIO.StringIO('')
  42
  43         d = fileSender.beginFileTransfer(emptyFile, consumer, lambda x: x)
  44
  45         # The producer will be immediately exhausted, and so immediately
  46         # unregistered
  47         self.assertEqual(consumer.producer, None)
  48
  49         # Which means the Deferred from FileSender should have been called
  50         self.failUnless(d.called,
  51                         'producer unregistered with deferred being called')
  52