<s> from twisted . trial import unittest <EOL> from twisted . internet import defer , reactor <EOL> from twisted . internet . endpoints import TCP4ServerEndpoint <EOL> from twisted . internet . protocol import ServerFactory <EOL> from twisted . protocols . basic import Int32StringReceiver <EOL> from tensor . ihateprotobuf import proto_pb2 <EOL> from tensor . objects import Event , Source , Output <EOL> from tensor . protocol . riemann import RiemannClientFactory <EOL> from tensor . service import TensorService <EOL> from tensor . aggregators import Counter32 , Counter64 , Counter <EOL> def wait ( secs ) : <EOL> d = defer . Deferred ( ) <EOL> reactor . callLater ( secs , d . callback , None ) <EOL> return d <EOL> class FakeRiemannServerProtocol ( Int32StringReceiver ) : <EOL> def stringReceived ( self , string ) : <EOL> msg = proto_pb2 . Msg . FromString ( string ) <EOL> self . factory . receive_message ( msg ) <EOL> self . sendString ( proto_pb2 . Msg ( ok = True ) . SerializeToString ( ) ) <EOL> class FakeRiemannServerFactory ( ServerFactory ) : <EOL> protocol = FakeRiemannServerProtocol <EOL> endpoint = None <EOL> listener = None <EOL> def __init__ ( self ) : <EOL> self . received_messages = [ ] <EOL> self . waiters = [ ] <EOL> def start_listening ( self , endpoint ) : <EOL> def cb ( listener ) : <EOL> self . listener = listener <EOL> return self <EOL> self . endpoint = endpoint <EOL> return endpoint . listen ( self ) . addCallback ( cb ) <EOL> def stop_listening ( self ) : <EOL> if self . listener is not None : <EOL> return self . listener . stopListening ( ) <EOL> def reconnect ( self ) : <EOL> return self . listener . startListening ( ) <EOL> def get_host ( self ) : <EOL> if self . listener is not None : <EOL> return self . listener . getHost ( ) <EOL> def receive_message ( self , msg ) : <EOL> self . received_messages . append ( msg ) <EOL> self . _process_waiters ( ) <EOL> def _process_waiters ( self ) : <EOL> new_waiters = [ ] <EOL> firing_waiters = [ ] <EOL> while self . waiters : <EOL> d , total_messages = self . waiters . pop ( <NUM_LIT:0> ) <EOL> if len ( self . received_messages ) >= total_messages : <EOL> firing_waiters . append ( d ) <EOL> else : <EOL> new_waiters . append ( ( d , total_messages ) ) <EOL> self . waiters = new_waiters <EOL> for d in firing_waiters : <EOL> d . callback ( self . received_messages [ : ] ) <EOL> def wait_for_messages ( self , total_messages ) : <EOL> d = defer . Deferred ( ) <EOL> self . waiters . append ( ( d , total_messages ) ) <EOL> self . _process_waiters ( ) <EOL> return d <EOL> class FakeSource ( Source ) : <EOL> pass <EOL> class FakeOutput ( Output ) : <EOL> def __init__ ( self , * a ) : <EOL> Output . __init__ ( self , * a ) <EOL> self . events = None <EOL> def eventsReceived ( self , events ) : <EOL> self . events = events <EOL> class TestService ( unittest . TestCase ) : <EOL> timeout = <NUM_LIT:10> <EOL> def start_riemann_server ( self ) : <EOL> factory = FakeRiemannServerFactory ( ) <EOL> self . addCleanup ( factory . stop_listening ) <EOL> return factory . start_listening ( TCP4ServerEndpoint ( reactor , <NUM_LIT:0> ) ) <EOL> def make_service ( self , config ) : <EOL> service = TensorService ( config ) <EOL> self . addCleanup ( self . stop_service , service ) <EOL> return service <EOL> def make_source ( self , service ) : <EOL> source = FakeSource ( <EOL> { <EOL> '<STR_LIT>' : '<STR_LIT:test>' , <EOL> '<STR_LIT>' : <NUM_LIT:1.0> , <EOL> '<STR_LIT>' : <NUM_LIT> <EOL> } , <EOL> service . sendEvent , service ) <EOL> service . sources . append ( source ) <EOL> return source <EOL> @ defer . inlineCallbacks <EOL> def stop_service ( self , service ) : <EOL> """<STR_LIT>""" <EOL> yield wait ( <NUM_LIT:0.1> ) <EOL> yield service . stopService ( ) <EOL> yield wait ( <NUM_LIT:0.1> ) <EOL> @ defer . inlineCallbacks <EOL> def test_service_sends_event ( self ) : <EOL> factory = yield self . start_riemann_server ( ) <EOL> service = self . make_service ( { "<STR_LIT:port>" : factory . get_host ( ) . port } ) <EOL> yield service . startService ( ) <EOL> [ ] = yield factory . wait_for_messages ( <NUM_LIT:0> ) <EOL> source = self . make_source ( service ) <EOL> event = Event ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <NUM_LIT:1.0> , <NUM_LIT> , <EOL> hostname = '<STR_LIT:localhost>' ) <EOL> service . sendEvent ( source , event ) <EOL> [ msg ] = yield factory . wait_for_messages ( <NUM_LIT:1> ) <EOL> [ event ] = msg . events <EOL> self . assertEqual ( event . description , '<STR_LIT>' ) <EOL> @ defer . inlineCallbacks <EOL> def test_service_sends_event_after_reconnect ( self ) : <EOL> self . patch ( RiemannClientFactory , '<STR_LIT>' , <NUM_LIT:0.1> ) <EOL> factory = yield self . start_riemann_server ( ) <EOL> service = self . make_service ( { "<STR_LIT:port>" : factory . get_host ( ) . port } ) <EOL> yield service . startService ( ) <EOL> [ ] = yield factory . wait_for_messages ( <NUM_LIT:0> ) <EOL> source = self . make_source ( service ) <EOL> event = Event ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <NUM_LIT:1.0> , <NUM_LIT> , <EOL> hostname = '<STR_LIT:localhost>' ) <EOL> service . sendEvent ( source , event ) <EOL> [ msg ] = yield factory . wait_for_messages ( <NUM_LIT:1> ) <EOL> [ event ] = msg . events <EOL> self . assertEqual ( event . description , '<STR_LIT>' ) <EOL> [ output ] = service . outputs [ None ] <EOL> yield output . connector . disconnect ( ) <EOL> yield wait ( <NUM_LIT> ) <EOL> [ _ ] = yield factory . wait_for_messages ( <NUM_LIT:1> ) <EOL> event = Event ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <NUM_LIT:1.0> , <NUM_LIT> , <EOL> hostname = '<STR_LIT:localhost>' ) <EOL> service . sendEvent ( source , event ) <EOL> [ _ , msg2 ] = yield factory . wait_for_messages ( <NUM_LIT:2> ) <EOL> [ event ] = msg . events <EOL> self . assertEqual ( event . description , '<STR_LIT>' ) <EOL> def _aggregator_test ( self , m1 , m2 , aggregator , delta ) : <EOL> service = self . make_service ( { } ) <EOL> ev1 = Event ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , m1 , delta , <EOL> hostname = '<STR_LIT:localhost>' , aggregation = aggregator ) <EOL> ev2 = Event ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , m2 , delta , <EOL> hostname = '<STR_LIT:localhost>' , aggregation = aggregator ) <EOL> ev1 . time = <NUM_LIT:1> <EOL> ev2 . time = delta + <NUM_LIT:1> <EOL> q1 = service . _aggregateQueue ( [ ev1 ] ) <EOL> ev = service . _aggregateQueue ( [ ev2 ] ) [ <NUM_LIT:0> ] <EOL> self . assertEqual ( q1 , [ ] ) <EOL> return ev . metric <EOL> def test_aggregate_counter32 ( self ) : <EOL> metric = self . _aggregator_test ( <NUM_LIT:1> , <NUM_LIT:2> , Counter32 , <NUM_LIT:4> ) <EOL> self . assertEqual ( metric , <NUM_LIT> ) <EOL> def test_aggregate_counter64 ( self ) : <EOL> metric = self . _aggregator_test ( <NUM_LIT:1> , <NUM_LIT:2> , Counter64 , <NUM_LIT:4> ) <EOL> self . assertEqual ( metric , <NUM_LIT> ) <EOL> def test_aggregate_counter ( self ) : <EOL> metric = self . _aggregator_test ( <NUM_LIT:1> , <NUM_LIT:2> , Counter , <NUM_LIT:4> ) <EOL> self . assertEqual ( metric , <NUM_LIT> ) <EOL> def test_aggregate_counter32_rollover ( self ) : <EOL> metric = self . _aggregator_test ( <NUM_LIT> , <NUM_LIT:5> , Counter32 , <NUM_LIT:4> ) <EOL> self . assertEqual ( metric , <NUM_LIT> ) <EOL> def test_aggregate_counter64_rollover ( self ) : <EOL> metric = self . _aggregator_test ( <NUM_LIT> , <NUM_LIT:5> , Counter64 , <NUM_LIT:4> ) <EOL> self . assertEqual ( metric , <NUM_LIT> ) <EOL> def test_state_match ( self ) : <EOL> service = self . make_service ( { <EOL> '<STR_LIT>' : <NUM_LIT:1.0> , '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : [ { <EOL> '<STR_LIT:source>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : { '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : { '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : '<STR_LIT>' } ] <EOL> } ) <EOL> ev1 = Event ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <NUM_LIT:50> , <NUM_LIT:1> , <EOL> hostname = '<STR_LIT:localhost>' ) <EOL> ev2 = Event ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <NUM_LIT:1000> , <NUM_LIT:1> , <EOL> hostname = '<STR_LIT:localhost>' ) <EOL> ev3 = Event ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <NUM_LIT:200> , <NUM_LIT:1> , <EOL> hostname = '<STR_LIT:localhost>' ) <EOL> service . setStates ( service . sources [ <NUM_LIT:0> ] , [ ev1 , ev2 , ev3 ] ) <EOL> self . assertEqual ( ev1 . state , '<STR_LIT>' ) <EOL> self . assertEqual ( ev2 . state , '<STR_LIT>' ) <EOL> self . assertEqual ( ev3 . state , '<STR_LIT>' ) <EOL> @ defer . inlineCallbacks <EOL> def test_source_routing ( self ) : <EOL> service = self . make_service ( { <EOL> '<STR_LIT>' : <NUM_LIT:1.0> , '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : [ { <EOL> '<STR_LIT:source>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : <NUM_LIT> , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' } ] <EOL> } ) <EOL> output1 = FakeOutput ( { } , service ) <EOL> output2 = FakeOutput ( { } , service ) <EOL> [ source ] = service . sources <EOL> service . outputs = { <EOL> '<STR_LIT>' : [ output1 ] , <EOL> '<STR_LIT>' : [ output2 ] <EOL> } <EOL> event = Event ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <NUM_LIT:1> , <NUM_LIT:1> , hostname = '<STR_LIT:localhost>' ) <EOL> service . sendEvent ( source , event ) <EOL> yield wait ( <NUM_LIT> ) <EOL> self . assertEqual ( len ( output1 . events ) , <NUM_LIT:1> ) <EOL> self . assertEqual ( output2 . events , None ) <EOL> output1 . events = None <EOL> source . config [ '<STR_LIT>' ] = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> service . sendEvent ( source , event ) <EOL> yield wait ( <NUM_LIT> ) <EOL> self . assertEqual ( len ( output1 . events ) , <NUM_LIT:1> ) <EOL> self . assertEqual ( len ( output2 . events ) , <NUM_LIT:1> ) </s>
