<s> import unittest <EOL> import gevent <EOL> from scales . constants import ChannelState <EOL> from scales . loadbalancer import HeapBalancerSink <EOL> from test . scales . loadbalancer . base_tests import LoadBalancerTestCase <EOL> class HeapBalancerTestCase ( LoadBalancerTestCase ) : <EOL> SINK_CLS = HeapBalancerSink <EOL> def testHeapOpensAllChannels ( self ) : <EOL> self . assertTrue ( all ( s . state == ChannelState . Open <EOL> for s in self . mock_provider . sinks_created ) ) <EOL> def testHeapForwardsMessage ( self ) : <EOL> self . _submitTestMessage ( ) <EOL> self . _completeTestMessage ( ) <EOL> self . assertIs ( self . return_message , self . MSG_SENTINEL ) <EOL> def testHeapMaintainsLeastLoad ( self ) : <EOL> hs = self . sink <EOL> self . _submitTestMessage ( ) <EOL> self . assertEqual ( hs . _heap [ <NUM_LIT:1> ] . load , hs . Idle ) <EOL> self . _completeTestMessage ( ) <EOL> self . assertEqual ( self . return_message , self . MSG_SENTINEL ) <EOL> def testHeapTracksLoad ( self ) : <EOL> self . _submitTestMessage ( ) <EOL> self . assertTrue ( any ( [ n . load == self . sink . Idle + <NUM_LIT:1> <EOL> for n in self . sink . _heap [ <NUM_LIT:1> : ] ] ) ) <EOL> self . _completeTestMessage ( ) <EOL> self . assertTrue ( all ( [ n . load == self . sink . Idle <EOL> for n in self . sink . _heap [ <NUM_LIT:1> : ] ] ) ) <EOL> def testHeapRemovesNode ( self ) : <EOL> self . _submitTestMessage ( ) <EOL> self . _completeTestMessage ( ) <EOL> self . mock_ss_provider . RemoveServer ( '<STR_LIT:localhost>' , <NUM_LIT> ) <EOL> self . assertTrue ( all ( n . channel . endpoint . port != <NUM_LIT> for n in self . sink . _heap [ <NUM_LIT:1> : ] ) ) <EOL> def testHeapRemovesLoadedNode ( self ) : <EOL> self . _submitTestMessage ( ) <EOL> loaded_node = self . _getLoadedNode ( ) <EOL> loaded_ep = loaded_node . channel . endpoint <EOL> self . mock_ss_provider . RemoveServer ( loaded_ep . host , loaded_ep . port ) <EOL> self . assertEqual ( loaded_node . index , - <NUM_LIT:1> ) <EOL> self . assertEqual ( loaded_node . load , self . sink . Idle + <NUM_LIT:1> ) <EOL> self . assertTrue ( all ( n . channel . endpoint . port != loaded_ep . port for n in self . sink . _heap [ <NUM_LIT:1> : ] ) ) <EOL> return loaded_node <EOL> def testHeapRemovedNodeReleaseCloses ( self ) : <EOL> loaded_node = self . testHeapRemovesLoadedNode ( ) <EOL> self . _completeTestMessage ( ) <EOL> self . assertEqual ( loaded_node . channel . state , ChannelState . Closed ) <EOL> def testHeapFailsOnEmpty ( self ) : <EOL> self . mock_ss_provider . RemoveAllServers ( ) <EOL> self . _submitTestMessage ( ) <EOL> self . assertIsNotNone ( self . return_message . error ) <EOL> def testHeapSucceedsWhenAllFailed ( self ) : <EOL> for s in self . mock_provider . sinks_created : <EOL> s . state = ChannelState . Closed <EOL> self . _submitTestMessage ( ) <EOL> self . assertIsNotNone ( self . return_message . error ) <EOL> def testHeapCorrectlyMarksNodesUp ( self ) : <EOL> self . testHeapSucceedsWhenAllFailed ( ) <EOL> self . mock_provider . sinks_created [ <NUM_LIT:0> ] . state = ChannelState . Open <EOL> self . _prepareSinkStack ( ) <EOL> self . _submitTestMessage ( ) <EOL> self . _completeTestMessage ( ) <EOL> self . assertEqual ( self . return_message , self . MSG_SENTINEL ) <EOL> def testHeapConcurrentOpenAndRequest ( self ) : <EOL> self . mock_ss_provider . RemoveAllServers ( ) <EOL> self . global_properties [ '<STR_LIT>' ] = <NUM_LIT> <EOL> def race ( ) : <EOL> gevent . sleep ( <NUM_LIT> ) <EOL> self . _submitTestMessage ( ) <EOL> self . _completeTestMessage ( ) <EOL> racer = gevent . spawn ( race ) <EOL> self . mock_ss_provider . AddServer ( '<STR_LIT:localhost>' , <NUM_LIT> ) <EOL> racer . join ( ) <EOL> self . assertEqual ( self . return_message , self . MSG_SENTINEL ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> unittest . main ( ) </s>
