<s> from struct import pack <EOL> import unittest <EOL> from cStringIO import StringIO <EOL> import gevent <EOL> from scales . thrift . sink import SocketTransportSink , ChannelConcurrencyError <EOL> from scales . constants import ChannelState <EOL> from test . scales . util . mocks import MockSocket <EOL> from test . scales . util . base import SinkTestCase <EOL> class ExpectedException ( Exception ) : pass <EOL> class ThriftSinkTestCase ( SinkTestCase ) : <EOL> SINK_CLS = SocketTransportSink <EOL> def _createSink ( self ) : <EOL> self . mock_socket = MockSocket ( '<STR_LIT:localhost>' , <NUM_LIT> ) <EOL> return SocketTransportSink ( self . mock_socket , '<STR_LIT:test>' ) <EOL> def _processTransaction ( self , request , response , open = None , close = None , read = None , write = None , sink = None , do_yield = True ) : <EOL> write_data = [ ] <EOL> read_data = [ pack ( '<STR_LIT>' , len ( response ) ) , response ] <EOL> def write_cb ( buff ) : <EOL> write_data . append ( buff ) <EOL> if write : <EOL> write ( buff ) <EOL> def read_cb ( sz ) : <EOL> if read : <EOL> read ( sz ) <EOL> return read_data . pop ( <NUM_LIT:0> ) <EOL> if not sink : <EOL> socket = MockSocket ( '<STR_LIT:localhost>' , <NUM_LIT> , open , close , read_cb , write_cb ) <EOL> sink = SocketTransportSink ( socket , '<STR_LIT:test>' ) <EOL> self . sink = sink <EOL> sink . Open ( ) . get ( ) <EOL> stream = StringIO ( request ) <EOL> self . _prepareSinkStack ( ) <EOL> sink_stack = self . sink_stack <EOL> sink . AsyncProcessRequest ( sink_stack , self . REQ_MSG_SENTINEL , stream , { } ) <EOL> if do_yield : <EOL> gevent . sleep ( ) <EOL> gevent . sleep ( ) <EOL> return read_data , write_data , sink , sink_stack <EOL> def testBasicTransaction ( self ) : <EOL> request = '<STR_LIT>' <EOL> response = '<STR_LIT>' <EOL> read_data , write_data , _ , _ = self . _processTransaction ( request , response ) <EOL> self . assertEqual ( write_data , [ pack ( '<STR_LIT>' , len ( request ) ) + request ] ) <EOL> self . assertEqual ( read_data , [ ] ) <EOL> self . assertEqual ( self . return_stream . getvalue ( ) , response ) <EOL> def testSocketOpenRaisesException ( self ) : <EOL> def open_cb ( ) : <EOL> raise ExpectedException ( ) <EOL> self . assertRaises ( ExpectedException , lambda : self . _processTransaction ( '<STR_LIT>' , '<STR_LIT>' , open_cb ) ) <EOL> self . assertEqual ( self . sink . state , ChannelState . Closed ) <EOL> def testConcurrency ( self ) : <EOL> request = '<STR_LIT>' <EOL> response = '<STR_LIT>' <EOL> _ , _ , sink , sink_stack = self . _processTransaction ( request , response , do_yield = False ) <EOL> _ , _ , _ , sink_stack_2 = self . _processTransaction ( request , response , sink = sink , do_yield = False ) <EOL> self . assertIsNotNone ( sink_stack_2 . return_message ) <EOL> self . assertIsInstance ( sink_stack_2 . return_message . error , ChannelConcurrencyError ) <EOL> gevent . sleep ( <NUM_LIT:0> ) <EOL> gevent . sleep ( <NUM_LIT:0> ) <EOL> self . assertEqual ( sink_stack . return_stream . getvalue ( ) , response ) <EOL> def testWriteFailure ( self ) : <EOL> def write_cb ( buff ) : <EOL> raise ExpectedException ( ) <EOL> self . _processTransaction ( '<STR_LIT:test>' , '<STR_LIT>' , write = write_cb ) <EOL> self . assertEqual ( self . sink . state , ChannelState . Closed ) <EOL> self . assertIsInstance ( self . sink_stack . return_message . error , ExpectedException ) <EOL> def testReadFailure ( self ) : <EOL> def read_cb ( sz ) : <EOL> raise ExpectedException ( ) <EOL> self . _processTransaction ( '<STR_LIT:test>' , '<STR_LIT>' , read = read_cb ) <EOL> self . assertEqual ( self . sink . state , ChannelState . Closed ) <EOL> self . assertIsInstance ( self . sink_stack . return_message . error , ExpectedException ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> unittest . main ( ) </s>
