<s> import collections <EOL> from gearman import compat <EOL> import unittest <EOL> from gearman . worker import GearmanWorker <EOL> from gearman . worker_handler import GearmanWorkerCommandHandler <EOL> from gearman . errors import ServerUnavailable , InvalidWorkerState <EOL> from gearman . protocol import get_command_name , GEARMAN_COMMAND_RESET_ABILITIES , GEARMAN_COMMAND_CAN_DO , GEARMAN_COMMAND_SET_CLIENT_ID , GEARMAN_COMMAND_NOOP , GEARMAN_COMMAND_PRE_SLEEP , GEARMAN_COMMAND_NO_JOB , GEARMAN_COMMAND_GRAB_JOB_UNIQ , GEARMAN_COMMAND_JOB_ASSIGN_UNIQ , GEARMAN_COMMAND_WORK_STATUS , GEARMAN_COMMAND_WORK_FAIL , GEARMAN_COMMAND_WORK_COMPLETE , GEARMAN_COMMAND_WORK_DATA , GEARMAN_COMMAND_WORK_EXCEPTION , GEARMAN_COMMAND_WORK_WARNING <EOL> from tests . _core_testing import _GearmanAbstractTest , MockGearmanConnectionManager , MockGearmanConnection <EOL> class MockGearmanWorker ( MockGearmanConnectionManager , GearmanWorker ) : <EOL> def __init__ ( self , * largs , ** kwargs ) : <EOL> super ( MockGearmanWorker , self ) . __init__ ( * largs , ** kwargs ) <EOL> self . worker_job_queues = compat . defaultdict ( collections . deque ) <EOL> def on_job_execute ( self , current_job ) : <EOL> current_handler = self . connection_to_handler_map [ current_job . connection ] <EOL> self . worker_job_queues [ current_handler ] . append ( current_job ) <EOL> class _GearmanAbstractWorkerTest ( _GearmanAbstractTest ) : <EOL> connection_manager_class = MockGearmanWorker <EOL> command_handler_class = GearmanWorkerCommandHandler <EOL> def setup_command_handler ( self ) : <EOL> super ( _GearmanAbstractWorkerTest , self ) . setup_command_handler ( ) <EOL> self . assert_sent_abilities ( [ ] ) <EOL> self . assert_sent_command ( GEARMAN_COMMAND_PRE_SLEEP ) <EOL> def assert_sent_abilities ( self , expected_abilities ) : <EOL> observed_abilities = set ( ) <EOL> self . assert_sent_command ( GEARMAN_COMMAND_RESET_ABILITIES ) <EOL> for ability in expected_abilities : <EOL> cmd_type , cmd_args = self . connection . _outgoing_commands . popleft ( ) <EOL> self . assertEqual ( get_command_name ( cmd_type ) , get_command_name ( GEARMAN_COMMAND_CAN_DO ) ) <EOL> observed_abilities . add ( cmd_args [ '<STR_LIT>' ] ) <EOL> self . assertEqual ( observed_abilities , set ( expected_abilities ) ) <EOL> def assert_sent_client_id ( self , expected_client_id ) : <EOL> self . assert_sent_command ( GEARMAN_COMMAND_SET_CLIENT_ID , client_id = expected_client_id ) <EOL> class WorkerTest ( _GearmanAbstractWorkerTest ) : <EOL> """<STR_LIT>""" <EOL> def test_registering_functions ( self ) : <EOL> def fake_callback_one ( worker_command_handler , current_job ) : <EOL> pass <EOL> def fake_callback_two ( worker_command_handler , current_job ) : <EOL> pass <EOL> self . connection_manager . register_task ( '<STR_LIT>' , fake_callback_one ) <EOL> self . failUnless ( '<STR_LIT>' in self . connection_manager . worker_abilities ) <EOL> self . failIf ( '<STR_LIT>' in self . connection_manager . worker_abilities ) <EOL> self . assertEqual ( self . connection_manager . worker_abilities [ '<STR_LIT>' ] , fake_callback_one ) <EOL> self . assertEqual ( self . command_handler . _handler_abilities , [ '<STR_LIT>' ] ) <EOL> self . connection_manager . register_task ( '<STR_LIT>' , fake_callback_two ) <EOL> self . failUnless ( '<STR_LIT>' in self . connection_manager . worker_abilities ) <EOL> self . failUnless ( '<STR_LIT>' in self . connection_manager . worker_abilities ) <EOL> self . assertEqual ( self . connection_manager . worker_abilities [ '<STR_LIT>' ] , fake_callback_one ) <EOL> self . assertEqual ( self . connection_manager . worker_abilities [ '<STR_LIT>' ] , fake_callback_two ) <EOL> self . assertEqual ( self . command_handler . _handler_abilities , [ '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> self . connection_manager . unregister_task ( '<STR_LIT>' ) <EOL> self . failIf ( '<STR_LIT>' in self . connection_manager . worker_abilities ) <EOL> self . failUnless ( '<STR_LIT>' in self . connection_manager . worker_abilities ) <EOL> self . assertEqual ( self . connection_manager . worker_abilities [ '<STR_LIT>' ] , fake_callback_two ) <EOL> self . assertEqual ( self . command_handler . _handler_abilities , [ '<STR_LIT>' ] ) <EOL> def test_setting_client_id ( self ) : <EOL> new_client_id = '<STR_LIT>' <EOL> self . assertEqual ( self . connection_manager . worker_client_id , None ) <EOL> self . assertEqual ( self . command_handler . _client_id , None ) <EOL> self . connection_manager . set_client_id ( new_client_id ) <EOL> self . assertEqual ( self . connection_manager . worker_client_id , new_client_id ) <EOL> self . assertEqual ( self . command_handler . _client_id , new_client_id ) <EOL> def test_establish_worker_connections ( self ) : <EOL> self . connection_manager . connection_list = [ ] <EOL> self . connection_manager . command_handlers = { } <EOL> good_connection = MockGearmanConnection ( ) <EOL> good_connection . connect ( ) <EOL> good_connection . _fail_on_bind = False <EOL> failed_then_retried_connection = MockGearmanConnection ( ) <EOL> failed_then_retried_connection . _fail_on_bind = False <EOL> failed_connection = MockGearmanConnection ( ) <EOL> failed_connection . _fail_on_bind = True <EOL> self . connection_manager . connection_list = [ good_connection , failed_then_retried_connection , failed_connection ] <EOL> alive_connections = self . connection_manager . establish_worker_connections ( ) <EOL> self . assertTrue ( good_connection in alive_connections ) <EOL> self . assertTrue ( failed_then_retried_connection in alive_connections ) <EOL> self . assertFalse ( failed_connection in alive_connections ) <EOL> def test_establish_worker_connections_dead ( self ) : <EOL> self . connection_manager . connection_list = [ ] <EOL> self . connection_manager . command_handlers = { } <EOL> self . assertRaises ( ServerUnavailable , self . connection_manager . work ) <EOL> dead_connection = MockGearmanConnection ( ) <EOL> dead_connection . _fail_on_bind = True <EOL> dead_connection . connected = False <EOL> self . connection_manager . connection_list = [ dead_connection ] <EOL> self . assertRaises ( ServerUnavailable , self . connection_manager . work ) <EOL> class WorkerCommandHandlerInterfaceTest ( _GearmanAbstractWorkerTest ) : <EOL> """<STR_LIT>""" <EOL> def test_on_connect ( self ) : <EOL> expected_abilities = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> expected_client_id = '<STR_LIT>' <EOL> self . connection . connected = False <EOL> self . connection_manager . set_client_id ( expected_client_id ) <EOL> self . connection_manager . unregister_task ( '<STR_LIT>' ) <EOL> for task in expected_abilities : <EOL> self . connection_manager . register_task ( task , None ) <EOL> self . connection_manager . establish_connection ( self . connection ) <EOL> self . assertNotEquals ( self . command_handler , self . connection_manager . connection_to_handler_map [ self . connection ] ) <EOL> self . assert_sent_client_id ( expected_client_id ) <EOL> self . assert_sent_abilities ( expected_abilities ) <EOL> self . assert_sent_command ( GEARMAN_COMMAND_PRE_SLEEP ) <EOL> self . assert_no_pending_commands ( ) <EOL> def test_set_abilities ( self ) : <EOL> expected_abilities = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> self . command_handler . set_abilities ( expected_abilities ) <EOL> self . assert_sent_abilities ( expected_abilities ) <EOL> self . assert_no_pending_commands ( ) <EOL> def test_set_client_id ( self ) : <EOL> expected_client_id = '<STR_LIT>' <EOL> handler_initial_state = { } <EOL> handler_initial_state [ '<STR_LIT>' ] = [ ] <EOL> handler_initial_state [ '<STR_LIT>' ] = None <EOL> self . command_handler . set_client_id ( expected_client_id ) <EOL> self . assert_sent_client_id ( expected_client_id ) <EOL> self . assert_no_pending_commands ( ) <EOL> def test_send_functions ( self ) : <EOL> current_job = self . generate_job ( ) <EOL> self . command_handler . send_job_status ( current_job , <NUM_LIT:0> , <NUM_LIT:1> ) <EOL> self . assert_sent_command ( GEARMAN_COMMAND_WORK_STATUS , job_handle = current_job . handle , numerator = '<STR_LIT:0>' , denominator = '<STR_LIT:1>' ) <EOL> self . command_handler . send_job_complete ( current_job , '<STR_LIT>' ) <EOL> self . assert_sent_command ( GEARMAN_COMMAND_WORK_COMPLETE , job_handle = current_job . handle , data = '<STR_LIT>' ) <EOL> self . command_handler . send_job_failure ( current_job ) <EOL> self . assert_sent_command ( GEARMAN_COMMAND_WORK_FAIL , job_handle = current_job . handle ) <EOL> self . command_handler . send_job_exception ( current_job , '<STR_LIT>' ) <EOL> self . assert_sent_command ( GEARMAN_COMMAND_WORK_EXCEPTION , job_handle = current_job . handle , data = '<STR_LIT>' ) <EOL> self . command_handler . send_job_data ( current_job , '<STR_LIT>' ) <EOL> self . assert_sent_command ( GEARMAN_COMMAND_WORK_DATA , job_handle = current_job . handle , data = '<STR_LIT>' ) <EOL> self . command_handler . send_job_warning ( current_job , '<STR_LIT>' ) <EOL> self . assert_sent_command ( GEARMAN_COMMAND_WORK_WARNING , job_handle = current_job . handle , data = '<STR_LIT>' ) <EOL> class WorkerCommandHandlerStateMachineTest ( _GearmanAbstractWorkerTest ) : <EOL> """<STR_LIT>""" <EOL> connection_manager_class = MockGearmanWorker <EOL> command_handler_class = GearmanWorkerCommandHandler <EOL> def setup_connection_manager ( self ) : <EOL> super ( WorkerCommandHandlerStateMachineTest , self ) . setup_connection_manager ( ) <EOL> self . connection_manager . register_task ( '<STR_LIT>' , None ) <EOL> def setup_command_handler ( self ) : <EOL> super ( _GearmanAbstractWorkerTest , self ) . setup_command_handler ( ) <EOL> self . assert_sent_abilities ( [ '<STR_LIT>' ] ) <EOL> self . assert_sent_command ( GEARMAN_COMMAND_PRE_SLEEP ) <EOL> def test_wakeup_work ( self ) : <EOL> self . move_to_state_wakeup ( ) <EOL> self . move_to_state_job_assign_uniq ( self . generate_job_dict ( ) ) <EOL> self . move_to_state_wakeup ( ) <EOL> self . move_to_state_no_job ( ) <EOL> def test_wakeup_sleep_wakup_work ( self ) : <EOL> self . move_to_state_wakeup ( ) <EOL> self . move_to_state_no_job ( ) <EOL> self . move_to_state_wakeup ( ) <EOL> self . move_to_state_job_assign_uniq ( self . generate_job_dict ( ) ) <EOL> self . move_to_state_wakeup ( ) <EOL> self . move_to_state_no_job ( ) <EOL> def test_multiple_wakeup_then_no_work ( self ) : <EOL> self . move_to_state_wakeup ( ) <EOL> for _ in range ( <NUM_LIT:5> ) : <EOL> self . command_handler . recv_command ( GEARMAN_COMMAND_NOOP ) <EOL> self . assert_job_lock ( is_locked = True ) <EOL> self . move_to_state_no_job ( ) <EOL> def test_multiple_work ( self ) : <EOL> self . move_to_state_wakeup ( ) <EOL> self . move_to_state_job_assign_uniq ( self . generate_job_dict ( ) ) <EOL> self . move_to_state_wakeup ( ) <EOL> self . move_to_state_job_assign_uniq ( self . generate_job_dict ( ) ) <EOL> self . move_to_state_wakeup ( ) <EOL> self . move_to_state_job_assign_uniq ( self . generate_job_dict ( ) ) <EOL> self . move_to_state_wakeup ( ) <EOL> self . move_to_state_no_job ( ) <EOL> def test_worker_already_locked ( self ) : <EOL> other_connection = MockGearmanConnection ( ) <EOL> self . connection_manager . connection_list . append ( other_connection ) <EOL> self . connection_manager . establish_connection ( other_connection ) <EOL> other_handler = self . connection_manager . connection_to_handler_map [ other_connection ] <EOL> other_handler . recv_command ( GEARMAN_COMMAND_NOOP ) <EOL> self . assertEqual ( self . connection_manager . command_handler_holding_job_lock , other_handler ) <EOL> self . assert_no_pending_commands ( ) <EOL> self . command_handler . recv_command ( GEARMAN_COMMAND_NOOP ) <EOL> self . assert_sent_command ( GEARMAN_COMMAND_PRE_SLEEP ) <EOL> self . assertEqual ( self . connection_manager . command_handler_holding_job_lock , other_handler ) <EOL> other_handler . recv_command ( GEARMAN_COMMAND_NO_JOB ) <EOL> self . assert_job_lock ( is_locked = False ) <EOL> self . move_to_state_wakeup ( ) <EOL> self . move_to_state_job_assign_uniq ( self . generate_job_dict ( ) ) <EOL> self . move_to_state_wakeup ( ) <EOL> self . move_to_state_no_job ( ) <EOL> def move_to_state_wakeup ( self ) : <EOL> self . assert_no_pending_commands ( ) <EOL> self . assert_job_lock ( is_locked = False ) <EOL> self . command_handler . recv_command ( GEARMAN_COMMAND_NOOP ) <EOL> def move_to_state_no_job ( self ) : <EOL> """<STR_LIT>""" <EOL> self . assert_awaiting_job ( ) <EOL> self . command_handler . recv_command ( GEARMAN_COMMAND_NO_JOB ) <EOL> self . assert_sent_command ( GEARMAN_COMMAND_PRE_SLEEP ) <EOL> self . assert_no_pending_commands ( ) <EOL> self . assert_job_lock ( is_locked = False ) <EOL> def move_to_state_job_assign_uniq ( self , fake_job ) : <EOL> """<STR_LIT>""" <EOL> self . assert_awaiting_job ( ) <EOL> self . command_handler . recv_command ( GEARMAN_COMMAND_JOB_ASSIGN_UNIQ , ** fake_job ) <EOL> current_job = self . connection_manager . worker_job_queues [ self . command_handler ] . popleft ( ) <EOL> self . assertEqual ( current_job . handle , fake_job [ '<STR_LIT>' ] ) <EOL> self . assertEqual ( current_job . task , fake_job [ '<STR_LIT>' ] ) <EOL> self . assertEqual ( current_job . unique , fake_job [ '<STR_LIT>' ] ) <EOL> self . assertEqual ( current_job . data , fake_job [ '<STR_LIT:data>' ] ) <EOL> self . assert_job_lock ( is_locked = False ) <EOL> self . assert_sent_command ( GEARMAN_COMMAND_PRE_SLEEP ) <EOL> def assert_awaiting_job ( self ) : <EOL> self . assert_sent_command ( GEARMAN_COMMAND_GRAB_JOB_UNIQ ) <EOL> self . assert_no_pending_commands ( ) <EOL> def assert_job_lock ( self , is_locked ) : <EOL> expected_value = ( is_locked and self . command_handler ) or None <EOL> self . assertEqual ( self . connection_manager . command_handler_holding_job_lock , expected_value ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> unittest . main ( ) </s>
