<s> import unittest <EOL> import time <EOL> import threading <EOL> import ufora . networking . MultiChannelListener as MultiChannelListener <EOL> import ufora . distributed . SharedState . Connections . TcpChannelFactory as TcpChannelFactory <EOL> import ufora . distributed . util . common as common <EOL> import ufora . native . CallbackScheduler as CallbackScheduler <EOL> callbackScheduler = CallbackScheduler . singletonForTesting ( ) <EOL> class MultiChannelListenerTest ( unittest . TestCase ) : <EOL> def setUp ( self ) : <EOL> self . defaultGroupId = <NUM_LIT> <EOL> self . ports = [ <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] <EOL> self . connectionCounts = { } <EOL> self . multiChannelListener = MultiChannelListener . MultiChannelListener ( callbackScheduler , self . ports ) <EOL> self . channelFactory = TcpChannelFactory . TcpStringChannelFactory ( callbackScheduler ) <EOL> def tearDown ( self ) : <EOL> self . multiChannelListener . stop ( ) <EOL> def createChannel ( self , port ) : <EOL> t0 = time . time ( ) <EOL> lastException = None <EOL> while t0 + <NUM_LIT> > time . time ( ) : <EOL> try : <EOL> return self . channelFactory . createChannel ( ( '<STR_LIT:localhost>' , port ) ) <EOL> except common . SocketException as e : <EOL> lastException = e <EOL> time . sleep ( <NUM_LIT> ) <EOL> assert lastException is not None <EOL> raise lastException <EOL> def startListener ( self ) : <EOL> self . multiChannelListener . start ( wait = True ) <EOL> self . connectionCounts = { ix : <NUM_LIT:0> for ix in range ( len ( self . ports ) ) } <EOL> def noopCompletionCallback ( self , channels , groupId ) : <EOL> self . assertEqual ( groupId , self . defaultGroupId ) <EOL> def noopConnectionCallback ( self , portIndex , channel ) : <EOL> self . multiChannelListener . setGroupIdForAcceptedChannel ( <EOL> channel , <EOL> self . defaultGroupId + self . connectionCounts [ portIndex ] <EOL> ) <EOL> self . connectionCounts [ portIndex ] += <NUM_LIT:1> <EOL> def test_init ( self ) : <EOL> self . assertEqual ( len ( self . multiChannelListener . listeners ) , len ( self . ports ) ) <EOL> def test_str ( self ) : <EOL> self . assertEqual ( <EOL> str ( self . multiChannelListener ) , <EOL> '<STR_LIT>' % '<STR_LIT:U+002CU+0020>' . join ( [ str ( p ) for p in self . ports ] ) <EOL> ) <EOL> def test_start_stop ( self ) : <EOL> self . multiChannelListener . registerConnectCallbackForAllPorts ( self . noopConnectionCallback ) <EOL> self . multiChannelListener . registerConnectionCompleteCallback ( self . noopCompletionCallback ) <EOL> self . startListener ( ) <EOL> self . multiChannelListener . stop ( ) <EOL> def test_register_callback_for_missing_port ( self ) : <EOL> with self . assertRaises ( Exception ) : <EOL> self . multiChannelListener . registerConnectCallbackForPort ( <NUM_LIT> , self . noopConnectionCallback ) <EOL> def test_register_callback_for_port ( self ) : <EOL> callbackFired = threading . Event ( ) <EOL> def onConnection ( portIndex , channel ) : <EOL> self . assertEqual ( portIndex , <NUM_LIT:0> ) <EOL> self . assertIsNotNone ( channel ) <EOL> callbackFired . set ( ) <EOL> return <NUM_LIT:1> <EOL> self . multiChannelListener . registerConnectCallbackForPort ( <NUM_LIT:0> , onConnection ) <EOL> for i in range ( <NUM_LIT:1> , len ( self . ports ) ) : <EOL> self . multiChannelListener . registerConnectCallbackForPort ( i , self . noopConnectionCallback ) <EOL> self . multiChannelListener . registerConnectionCompleteCallback ( self . noopCompletionCallback ) <EOL> self . startListener ( ) <EOL> channel = self . createChannel ( self . multiChannelListener . listeners [ <NUM_LIT:0> ] . port ) <EOL> self . assertIsNotNone ( channel ) <EOL> self . assertTrue ( callbackFired . wait ( <NUM_LIT:0.1> ) ) <EOL> channel . disconnect ( ) <EOL> def test_connection_completion ( self ) : <EOL> callbackFired = threading . Event ( ) <EOL> def onCompletion ( channels , groupId ) : <EOL> self . assertEqual ( len ( channels ) , len ( self . ports ) ) <EOL> self . assertEqual ( groupId , self . defaultGroupId ) <EOL> callbackFired . set ( ) <EOL> self . multiChannelListener . registerConnectionCompleteCallback ( onCompletion ) <EOL> self . multiChannelListener . registerConnectCallbackForAllPorts ( self . noopConnectionCallback ) <EOL> self . startListener ( ) <EOL> self . multiChannelListener . blockUntilReady ( ) <EOL> clients = [ self . createChannel ( listener . port ) for listener in self . multiChannelListener . listeners ] <EOL> self . assertNotIn ( None , clients ) <EOL> self . assertTrue ( callbackFired . wait ( <NUM_LIT:0.1> ) ) <EOL> for client in clients : <EOL> client . disconnect ( ) <EOL> def test_multiple_channel_groups ( self ) : <EOL> done = threading . Event ( ) <EOL> seenChannels = { ix : [ ] for ix in range ( len ( self . ports ) ) } <EOL> remainingGroups = [ <NUM_LIT:3> ] <EOL> def onConnection ( portIndex , channel ) : <EOL> seenChannels [ portIndex ] . append ( channel ) <EOL> self . multiChannelListener . setGroupIdForAcceptedChannel ( channel , len ( seenChannels [ portIndex ] ) ) <EOL> def onCompletion ( channels , groupId ) : <EOL> remainingGroups [ <NUM_LIT:0> ] -= <NUM_LIT:1> <EOL> if ( remainingGroups [ <NUM_LIT:0> ] == <NUM_LIT:0> ) : <EOL> done . set ( ) <EOL> self . multiChannelListener . registerConnectCallbackForAllPorts ( onConnection ) <EOL> self . multiChannelListener . registerConnectionCompleteCallback ( onCompletion ) <EOL> self . startListener ( ) <EOL> clients = { listener . index : [ self . createChannel ( listener . port ) for i in range ( remainingGroups [ <NUM_LIT:0> ] ) ] for listener in self . multiChannelListener . listeners } <EOL> self . assertTrue ( done . wait ( <NUM_LIT:0.1> ) ) <EOL> for channels in clients . itervalues ( ) : <EOL> for channel in channels : <EOL> channel . disconnect ( ) <EOL> def test_orphan_channel_group ( self ) : <EOL> def onCompletion ( channels , groupId ) : <EOL> self . assertTrue ( False ) <EOL> groupCompletionTimeout = <NUM_LIT> <EOL> self . multiChannelListener . GROUP_COMPLETION_TIMEOUT_IN_SECONDS = groupCompletionTimeout <EOL> self . multiChannelListener . registerConnectCallbackForAllPorts ( self . noopConnectionCallback ) <EOL> self . multiChannelListener . registerConnectionCompleteCallback ( onCompletion ) <EOL> self . startListener ( ) <EOL> clients = [ self . createChannel ( self . multiChannelListener . ports [ i ] ) for i in range ( <NUM_LIT:1> , len ( self . ports ) ) ] <EOL> time . sleep ( groupCompletionTimeout * <NUM_LIT:2> ) <EOL> clients . append ( self . createChannel ( self . multiChannelListener . ports [ <NUM_LIT:1> ] ) ) <EOL> clients . append ( self . createChannel ( self . multiChannelListener . ports [ <NUM_LIT:0> ] ) ) <EOL> time . sleep ( groupCompletionTimeout ) <EOL> def test_accept_timeout ( self ) : <EOL> def onConnection ( portIndex , channel ) : <EOL> pass <EOL> acceptTimeout = <NUM_LIT> <EOL> self . multiChannelListener . ACCEPT_TIMEOUT_IN_SECONDS = acceptTimeout <EOL> self . multiChannelListener . registerConnectCallbackForAllPorts ( onConnection ) <EOL> self . multiChannelListener . registerConnectionCompleteCallback ( self . noopCompletionCallback ) <EOL> self . startListener ( ) <EOL> channels = [ self . createChannel ( self . multiChannelListener . listeners [ <NUM_LIT:0> ] . port ) . makeQueuelike ( callbackScheduler ) ] <EOL> time . sleep ( acceptTimeout * <NUM_LIT> ) <EOL> channels . append ( self . createChannel ( self . multiChannelListener . listeners [ <NUM_LIT:0> ] . port ) ) <EOL> with self . assertRaises ( UserWarning ) : <EOL> channels [ <NUM_LIT:0> ] . getTimeout ( <NUM_LIT:1.0> ) <EOL> for channel in channels : <EOL> channel . disconnect ( ) </s>
