<s> import logging <EOL> from logging import StreamHandler <EOL> import sys <EOL> import unittest <EOL> import time <EOL> import zmq <EOL> from zmq import Socket <EOL> from zmq . eventloop . ioloop import IOLoop <EOL> from zmq . eventloop . zmqstream import ZMQStream <EOL> from spyder . core . constants import ZMQ_SPYDER_MGMT_WORKER <EOL> from spyder . core . constants import ZMQ_SPYDER_MGMT_WORKER_QUIT <EOL> from spyder . core . constants import ZMQ_SPYDER_MGMT_WORKER_QUIT_ACK <EOL> from spyder . core . mgmt import ZmqMgmt <EOL> from spyder . core . worker import ZmqWorker , AsyncZmqWorker <EOL> from spyder . core . messages import DataMessage , MgmtMessage <EOL> from spyder . thrift . gen . ttypes import CrawlUri <EOL> class ZmqTornadoIntegrationTest ( unittest . TestCase ) : <EOL> def setUp ( self ) : <EOL> self . _io_loop = IOLoop . instance ( ) <EOL> self . _ctx = zmq . Context ( <NUM_LIT:1> ) <EOL> self . _setup_mgmt_sockets ( ) <EOL> self . _setup_data_sockets ( ) <EOL> self . _mgmt = ZmqMgmt ( self . _mgmt_sockets [ '<STR_LIT>' ] , <EOL> self . _mgmt_sockets [ '<STR_LIT>' ] , io_loop = self . _io_loop ) <EOL> self . _mgmt . start ( ) <EOL> self . _mgmt . add_callback ( ZMQ_SPYDER_MGMT_WORKER , self . on_mgmt_end ) <EOL> def tearDown ( self ) : <EOL> self . _mgmt . stop ( ) <EOL> for socket in self . _mgmt_sockets . itervalues ( ) : <EOL> socket . close ( ) <EOL> for socket in self . _worker_sockets . itervalues ( ) : <EOL> socket . close ( ) <EOL> self . _ctx . term ( ) <EOL> def _setup_mgmt_sockets ( self ) : <EOL> self . _mgmt_sockets = dict ( ) <EOL> mgmt_master_worker = '<STR_LIT>' <EOL> sock = self . _ctx . socket ( zmq . PUB ) <EOL> sock . bind ( mgmt_master_worker ) <EOL> self . _mgmt_sockets [ '<STR_LIT>' ] = sock <EOL> self . _mgmt_sockets [ '<STR_LIT>' ] = ZMQStream ( sock , self . _io_loop ) <EOL> self . _mgmt_sockets [ '<STR_LIT>' ] = self . _ctx . socket ( zmq . SUB ) <EOL> self . _mgmt_sockets [ '<STR_LIT>' ] . setsockopt ( zmq . SUBSCRIBE , "<STR_LIT>" ) <EOL> self . _mgmt_sockets [ '<STR_LIT>' ] . connect ( mgmt_master_worker ) <EOL> mgmt_worker_master = '<STR_LIT>' <EOL> self . _mgmt_sockets [ '<STR_LIT>' ] = self . _ctx . socket ( zmq . PUB ) <EOL> self . _mgmt_sockets [ '<STR_LIT>' ] . bind ( mgmt_worker_master ) <EOL> sock = self . _ctx . socket ( zmq . SUB ) <EOL> sock . setsockopt ( zmq . SUBSCRIBE , "<STR_LIT>" ) <EOL> sock . connect ( mgmt_worker_master ) <EOL> self . _mgmt_sockets [ '<STR_LIT>' ] = sock <EOL> self . _mgmt_sockets [ '<STR_LIT>' ] = ZMQStream ( sock , self . _io_loop ) <EOL> def _setup_data_sockets ( self ) : <EOL> self . _worker_sockets = dict ( ) <EOL> data_master_worker = '<STR_LIT>' <EOL> sock = self . _ctx . socket ( zmq . PUSH ) <EOL> sock . bind ( data_master_worker ) <EOL> self . _worker_sockets [ '<STR_LIT>' ] = sock <EOL> self . _worker_sockets [ '<STR_LIT>' ] = ZMQStream ( sock , self . _io_loop ) <EOL> self . _worker_sockets [ '<STR_LIT>' ] = self . _ctx . socket ( zmq . PULL ) <EOL> self . _worker_sockets [ '<STR_LIT>' ] . connect ( data_master_worker ) <EOL> data_worker_master = '<STR_LIT>' <EOL> self . _worker_sockets [ '<STR_LIT>' ] = self . _ctx . socket ( zmq . PUB ) <EOL> self . _worker_sockets [ '<STR_LIT>' ] . bind ( data_worker_master ) <EOL> sock = self . _ctx . socket ( zmq . SUB ) <EOL> sock . setsockopt ( zmq . SUBSCRIBE , "<STR_LIT>" ) <EOL> sock . connect ( data_worker_master ) <EOL> self . _worker_sockets [ '<STR_LIT>' ] = sock <EOL> self . _worker_sockets [ '<STR_LIT>' ] = ZMQStream ( sock , self . _io_loop ) <EOL> def on_mgmt_end ( self , _msg ) : <EOL> self . _io_loop . stop ( ) <EOL> class ZmqWorkerIntegrationTest ( ZmqTornadoIntegrationTest ) : <EOL> def echo_processing ( self , crawl_uri ) : <EOL> death = MgmtMessage ( topic = ZMQ_SPYDER_MGMT_WORKER , <EOL> data = ZMQ_SPYDER_MGMT_WORKER_QUIT ) <EOL> self . _mgmt_sockets [ '<STR_LIT>' ] . send_multipart ( death . serialize ( ) ) <EOL> return crawl_uri <EOL> def test_that_stopping_worker_via_mgmt_works ( self ) : <EOL> worker = ZmqWorker ( self . _worker_sockets [ '<STR_LIT>' ] , <EOL> self . _worker_sockets [ '<STR_LIT>' ] , <EOL> self . _mgmt , <EOL> self . echo_processing , <EOL> StreamHandler ( sys . stdout ) , <EOL> logging . DEBUG , <EOL> self . _io_loop ) <EOL> worker . start ( ) <EOL> curi = CrawlUri ( url = "<STR_LIT>" ) <EOL> msg = DataMessage ( ) <EOL> msg . identity = "<STR_LIT>" <EOL> msg . curi = curi <EOL> def assert_correct_data_answer ( msg2 ) : <EOL> self . assertEqual ( msg , DataMessage ( msg2 ) ) <EOL> self . _worker_sockets [ '<STR_LIT>' ] . on_recv ( assert_correct_data_answer ) <EOL> def assert_correct_mgmt_answer ( msg3 ) : <EOL> self . assertEqual ( ZMQ_SPYDER_MGMT_WORKER_QUIT_ACK , msg3 . data ) <EOL> self . _mgmt_sockets [ '<STR_LIT>' ] . on_recv ( assert_correct_data_answer ) <EOL> self . _worker_sockets [ '<STR_LIT>' ] . send_multipart ( msg . serialize ( ) ) <EOL> self . _io_loop . start ( ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> unittest . main ( ) </s>
