test/py/TestEof.py

   1 #!/usr/bin/env python
   2
   3 import sys, glob
   4 sys.path.insert(0, './gen-py')
   5 sys.path.insert(0, glob.glob('../../lib/py/build/lib.*')[0])
   6
   7 from ThriftTest import ThriftTest
   8 from ThriftTest.ttypes import *
   9 from thrift.transport import TTransport
  10 from thrift.transport import TSocket
  11 from thrift.protocol import TBinaryProtocol
  12 import unittest
  13 import time
  14
  15 class TestEof(unittest.TestCase):
  16
  17   def setUp(self):
  18     trans = TTransport.TMemoryBuffer()
  19     prot = TBinaryProtocol.TBinaryProtocol(trans)
  20
  21     x = Xtruct()
  22     x.string_thing = "Zero"
  23     x.byte_thing = 0
  24
  25     x.write(prot)
  26
  27     x = Xtruct()
  28     x.string_thing = "One"
  29     x.byte_thing = 1
  30
  31     x.write(prot)
  32
  33     self.data = trans.getvalue()
  34
  35   def testTransportReadAll(self):
  36     """Test that readAll on any type of transport throws an EOFError"""
  37     trans = TTransport.TMemoryBuffer(self.data)
  38     trans.readAll(1)
  39
  40     try:
  41       trans.readAll(10000)
  42     except EOFError:
  43       return
  44
  45     self.fail("Should have gotten EOFError")
  46
  47   def eofTestHelper(self, pfactory):
  48     trans = TTransport.TMemoryBuffer(self.data)
  49     prot = pfactory.getProtocol(trans)
  50
  51     x = Xtruct()
  52     x.read(prot)
  53     self.assertEqual(x.string_thing, "Zero")
  54     self.assertEqual(x.byte_thing, 0)
  55
  56     x = Xtruct()
  57     x.read(prot)
  58     self.assertEqual(x.string_thing, "One")
  59     self.assertEqual(x.byte_thing, 1)
  60
  61     try:
  62       x = Xtruct()
  63       x.read(prot)
  64     except EOFError:
  65       return
  66
  67     self.fail("Should have gotten EOFError")
  68
  69   def eofTestHelperStress(self, pfactory):
  70     """Teest the ability of TBinaryProtocol to deal with the removal of every byte in the file"""
  71     # TODO: we should make sure this covers more of the code paths
  72
  73     for i in xrange(0, len(self.data) + 1):
  74       trans = TTransport.TMemoryBuffer(self.data[0:i])
  75       prot = pfactory.getProtocol(trans)
  76       try:
  77         x = Xtruct()
  78         x.read(prot)
  79         x.read(prot)
  80         x.read(prot)
  81       except EOFError:
  82         continue
  83       self.fail("Should have gotten an EOFError")
  84
  85   def testBinaryProtocolEof(self):
  86     """Test that TBinaryProtocol throws an EOFError when it reaches the end of the stream"""
  87     self.eofTestHelper(TBinaryProtocol.TBinaryProtocolFactory())
  88     self.eofTestHelperStress(TBinaryProtocol.TBinaryProtocolFactory())
  89
  90   def testBinaryProtocolAcceleratedEof(self):
  91     """Test that TBinaryProtocolAccelerated throws an EOFError when it reaches the end of the stream"""
  92     self.eofTestHelper(TBinaryProtocol.TBinaryProtocolAcceleratedFactory())
  93     self.eofTestHelperStress(TBinaryProtocol.TBinaryProtocolAcceleratedFactory())
  94
  95 suite = unittest.TestSuite()
  96 loader = unittest.TestLoader()
  97
  98 suite.addTest(loader.loadTestsFromTestCase(TestEof))
  99
 100 testRunner = unittest.TextTestRunner(verbosity=2)
 101 testRunner.run(suite)