import collections import random import unittest from gearman.client import GearmanClient from gearman.client_handler import GearmanClientCommandHandler from gearman.constants import PRIORITY_NONE, PRIORITY_HIGH, PRIORITY_LOW, JOB_UNKNOWN, JOB_PENDING, JOB_CREATED, JOB_FAILED, JOB_COMPLETE from gearman.errors import ExceededConnectionAttempts, ServerUnavailable, InvalidClientState from gearman.protocol import submit_cmd_for_background_priority, GEARMAN_COMMAND_STATUS_RES, GEARMAN_COMMAND_GET_STATUS, GEARMAN_COMMAND_JOB_CREATED, \ GEARMAN_COMMAND_WORK_STATUS, GEARMAN_COMMAND_WORK_FAIL, GEARMAN_COMMAND_WORK_COMPLETE, GEARMAN_COMMAND_WORK_DATA, GEARMAN_COMMAND_WORK_WARNING from tests._core_testing import _GearmanAbstractTest, MockGearmanConnectionManager, MockGearmanConnection class MockGearmanClient(GearmanClient, MockGearmanConnectionManager): pass class ClientTest(_GearmanAbstractTest): """Test the public client interface""" connection_manager_class = MockGearmanClient command_handler_class = GearmanClientCommandHandler def setUp(self): super(ClientTest, self).setUp() self.original_handle_connection_activity = self.connection_manager.handle_connection_activity def tearDown(self): super(ClientTest, self).tearDown() self.connection_manager.handle_connection_activity = self.original_handle_connection_activity def generate_job_request(self, submitted=True, accepted=True): current_request = super(ClientTest, self).generate_job_request() if submitted or accepted: self.connection_manager.establish_request_connection(current_request) self.command_handler.send_job_request(current_request) if submitted and accepted: self.command_handler.recv_command(GEARMAN_COMMAND_JOB_CREATED, job_handle=current_request.job.handle) self.assert_(current_request.job.handle in self.command_handler.handle_to_request_map) return current_request def test_establish_request_connection_complex(self): # Spin up a bunch of imaginary gearman connections failed_connection = MockGearmanConnection() failed_connection._fail_on_bind = True failed_then_retried_connection = MockGearmanConnection() failed_then_retried_connection._fail_on_bind = True good_connection = MockGearmanConnection() good_connection.connect() # Register all our connections self.connection_manager.connection_list = [failed_connection, failed_then_retried_connection, good_connection] # When we first create our request, our client shouldn't know anything about it current_request = self.generate_job_request(submitted=False, accepted=False) self.failIf(current_request in self.connection_manager.request_to_rotating_connection_queue) # Make sure that when we start up, we get our good connection chosen_connection = self.connection_manager.establish_request_connection(current_request) self.assertEqual(chosen_connection, good_connection) self.assertFalse(failed_connection.connected) self.assertFalse(failed_then_retried_connection.connected) self.assertTrue(good_connection.connected) # No state changed so we should still go to the correct connection chosen_connection = self.connection_manager.establish_request_connection(current_request) self.assertEqual(chosen_connection, good_connection) # Pretend like our good connection died so we'll need to choose somethign else good_connection._reset_connection() good_connection._fail_on_bind = True failed_then_retried_connection._fail_on_bind = False failed_then_retried_connection.connect() # Make sure we rotate good_connection and failed_connection out chosen_connection = self.connection_manager.establish_request_connection(current_request) self.assertEqual(chosen_connection, failed_then_retried_connection) self.assertFalse(failed_connection.connected) self.assertTrue(failed_then_retried_connection.connected) self.assertFalse(good_connection.connected) def test_establish_request_connection_dead(self): self.connection_manager.connection_list = [] self.connection_manager.command_handlers = {} current_request = self.generate_job_request(submitted=False, accepted=False) # No connections == death self.assertRaises(ServerUnavailable, self.connection_manager.establish_request_connection, current_request) # Spin up a bunch of imaginary gearman connections failed_connection = MockGearmanConnection() failed_connection._fail_on_bind = True self.connection_manager.connection_list.append(failed_connection) # All failed connections == death self.assertRaises(ServerUnavailable, self.connection_manager.establish_request_connection, current_request) def test_auto_retry_behavior(self): current_request = self.generate_job_request(submitted=False, accepted=False) def fail_then_create_jobs(rx_conns, wr_conns, ex_conns): if self.connection_manager.current_failures < self.connection_manager.expected_failures: self.connection_manager.current_failures += 1 # We're going to down this connection and reset state self.assertTrue(self.connection.connected) self.connection_manager.handle_error(self.connection) self.assertFalse(self.connection.connected) # We're then going to IMMEDIATELY pull this connection back up # So we don't bail out of the "self.connection_manager.poll_connections_until_stopped" loop self.connection_manager.establish_connection(self.connection) else: self.assertEquals(current_request.state, JOB_PENDING) self.command_handler.recv_command(GEARMAN_COMMAND_JOB_CREATED, job_handle=current_request.job.handle) return rx_conns, wr_conns, ex_conns self.connection_manager.handle_connection_activity = fail_then_create_jobs self.connection_manager.expected_failures = 5 # Now that we've setup our rety behavior, we need to reset the entire state of our experiment # First pass should succeed as we JUST touch our max attempts self.connection_manager.current_failures = current_request.connection_attempts = 0 current_request.max_connection_attempts = self.connection_manager.expected_failures + 1 current_request.state = JOB_UNKNOWN accepted_jobs = self.connection_manager.wait_until_jobs_accepted([current_request]) self.assertEquals(current_request.state, JOB_CREATED) self.assertEquals(current_request.connection_attempts, current_request.max_connection_attempts) # Second pass should fail as we JUST exceed our max attempts self.connection_manager.current_failures = current_request.connection_attempts = 0 current_request.max_connection_attempts = self.connection_manager.expected_failures current_request.state = JOB_UNKNOWN self.assertRaises(ExceededConnectionAttempts, self.connection_manager.wait_until_jobs_accepted, [current_request]) self.assertEquals(current_request.state, JOB_UNKNOWN) self.assertEquals(current_request.connection_attempts, current_request.max_connection_attempts) def test_multiple_fg_job_submission(self): submitted_job_count = 5 expected_job_list = [self.generate_job() for _ in xrange(submitted_job_count)] def mark_jobs_created(rx_conns, wr_conns, ex_conns): for current_job in expected_job_list: self.command_handler.recv_command(GEARMAN_COMMAND_JOB_CREATED, job_handle=current_job.handle) return rx_conns, wr_conns, ex_conns self.connection_manager.handle_connection_activity = mark_jobs_created job_dictionaries = [current_job.to_dict() for current_job in expected_job_list] # Test multiple job submission job_requests = self.connection_manager.submit_multiple_jobs(job_dictionaries, wait_until_complete=False) for current_request, expected_job in zip(job_requests, expected_job_list): current_job = current_request.job self.assert_jobs_equal(current_job, expected_job) self.assertEqual(current_request.priority, PRIORITY_NONE) self.assertEqual(current_request.background, False) self.assertEqual(current_request.state, JOB_CREATED) self.assertFalse(current_request.complete) def test_single_bg_job_submission(self): expected_job = self.generate_job() def mark_job_created(rx_conns, wr_conns, ex_conns): self.command_handler.recv_command(GEARMAN_COMMAND_JOB_CREATED, job_handle=expected_job.handle) return rx_conns, wr_conns, ex_conns self.connection_manager.handle_connection_activity = mark_job_created job_request = self.connection_manager.submit_job(expected_job.task, expected_job.data, unique=expected_job.unique, background=True, priority=PRIORITY_LOW, wait_until_complete=False) current_job = job_request.job self.assert_jobs_equal(current_job, expected_job) self.assertEqual(job_request.priority, PRIORITY_LOW) self.assertEqual(job_request.background, True) self.assertEqual(job_request.state, JOB_CREATED) self.assertTrue(job_request.complete) def test_single_fg_job_submission_timeout(self): expected_job = self.generate_job() def job_failed_submission(rx_conns, wr_conns, ex_conns): return rx_conns, wr_conns, ex_conns self.connection_manager.handle_connection_activity = job_failed_submission job_request = self.connection_manager.submit_job(expected_job.task, expected_job.data, unique=expected_job.unique, priority=PRIORITY_HIGH, poll_timeout=0.01) self.assertEqual(job_request.priority, PRIORITY_HIGH) self.assertEqual(job_request.background, False) self.assertEqual(job_request.state, JOB_PENDING) self.assertFalse(job_request.complete) self.assertTrue(job_request.timed_out) def test_wait_for_multiple_jobs_to_complete_or_timeout(self): completed_request = self.generate_job_request() failed_request = self.generate_job_request() timeout_request = self.generate_job_request() self.update_requests = True def multiple_job_updates(rx_conns, wr_conns, ex_conns): # Only give a single status update and have the 3rd job handle timeout if self.update_requests: self.command_handler.recv_command(GEARMAN_COMMAND_WORK_COMPLETE, job_handle=completed_request.job.handle, data='12345') self.command_handler.recv_command(GEARMAN_COMMAND_WORK_FAIL, job_handle=failed_request.job.handle) self.update_requests = False return rx_conns, wr_conns, ex_conns self.connection_manager.handle_connection_activity = multiple_job_updates finished_requests = self.connection_manager.wait_until_jobs_completed([completed_request, failed_request, timeout_request], poll_timeout=0.01) del self.update_requests finished_completed_request, finished_failed_request, finished_timeout_request = finished_requests self.assert_jobs_equal(finished_completed_request.job, completed_request.job) self.assertEqual(finished_completed_request.state, JOB_COMPLETE) self.assertEqual(finished_completed_request.result, '12345') self.assertFalse(finished_completed_request.timed_out) self.assert_(finished_completed_request.job.handle not in self.command_handler.handle_to_request_map) self.assert_jobs_equal(finished_failed_request.job, failed_request.job) self.assertEqual(finished_failed_request.state, JOB_FAILED) self.assertEqual(finished_failed_request.result, None) self.assertFalse(finished_failed_request.timed_out) self.assert_(finished_failed_request.job.handle not in self.command_handler.handle_to_request_map) self.assertEqual(finished_timeout_request.state, JOB_CREATED) self.assertEqual(finished_timeout_request.result, None) self.assertTrue(finished_timeout_request.timed_out) self.assert_(finished_timeout_request.job.handle in self.command_handler.handle_to_request_map) def test_get_job_status(self): single_request = self.generate_job_request() def retrieve_status(rx_conns, wr_conns, ex_conns): self.command_handler.recv_command(GEARMAN_COMMAND_STATUS_RES, job_handle=single_request.job.handle, known='1', running='0', numerator='0', denominator='1') return rx_conns, wr_conns, ex_conns self.connection_manager.handle_connection_activity = retrieve_status job_request = self.connection_manager.get_job_status(single_request) request_status = job_request.status self.failUnless(request_status) self.assertTrue(request_status['known']) self.assertFalse(request_status['running']) self.assertEqual(request_status['numerator'], 0) self.assertEqual(request_status['denominator'], 1) self.assertFalse(job_request.timed_out) def test_get_job_status_unknown(self): single_request = self.generate_job_request() current_handle = single_request.job.handle self.command_handler.recv_command(GEARMAN_COMMAND_WORK_FAIL, job_handle=current_handle) def retrieve_status(rx_conns, wr_conns, ex_conns): self.command_handler.recv_command(GEARMAN_COMMAND_STATUS_RES, job_handle=current_handle, known='0', running='0', numerator='0', denominator='1') return rx_conns, wr_conns, ex_conns self.connection_manager.handle_connection_activity = retrieve_status job_request = self.connection_manager.get_job_status(single_request) request_status = job_request.status self.failUnless(request_status) self.assertFalse(request_status['known']) self.assertFalse(request_status['running']) self.assertEqual(request_status['numerator'], 0) self.assertEqual(request_status['denominator'], 1) self.assertFalse(job_request.timed_out) self.assert_(current_handle not in self.command_handler.handle_to_request_map) def test_get_job_status_timeout(self): single_request = self.generate_job_request() def retrieve_status_timeout(rx_conns, wr_conns, ex_conns): pass self.connection_manager.handle_connection_activity = retrieve_status_timeout job_request = self.connection_manager.get_job_status(single_request, poll_timeout=0.01) self.assertTrue(job_request.timed_out) class ClientCommandHandlerInterfaceTest(_GearmanAbstractTest): """Test the public interface a GearmanClient may need to call in order to update state on a GearmanClientCommandHandler""" connection_manager_class = MockGearmanClient command_handler_class = GearmanClientCommandHandler def test_send_job_request(self): current_request = self.generate_job_request() gearman_job = current_request.job for priority in (PRIORITY_NONE, PRIORITY_HIGH, PRIORITY_LOW): for background in (False, True): current_request.reset() current_request.priority = priority current_request.background = background self.command_handler.send_job_request(current_request) queued_request = self.command_handler.requests_awaiting_handles.popleft() self.assertEqual(queued_request, current_request) expected_cmd_type = submit_cmd_for_background_priority(background, priority) self.assert_sent_command(expected_cmd_type, task=gearman_job.task, data=gearman_job.data, unique=gearman_job.unique) def test_get_status_of_job(self): current_request = self.generate_job_request() self.command_handler.send_get_status_of_job(current_request) self.assert_sent_command(GEARMAN_COMMAND_GET_STATUS, job_handle=current_request.job.handle) class ClientCommandHandlerStateMachineTest(_GearmanAbstractTest): """Test single state transitions within a GearmanWorkerCommandHandler""" connection_manager_class = MockGearmanClient command_handler_class = GearmanClientCommandHandler def generate_job_request(self, submitted=True, accepted=True): current_request = super(ClientCommandHandlerStateMachineTest, self).generate_job_request() if submitted or accepted: self.command_handler.requests_awaiting_handles.append(current_request) current_request.state = JOB_PENDING if submitted and accepted: self.command_handler.recv_command(GEARMAN_COMMAND_JOB_CREATED, job_handle=current_request.job.handle) return current_request def test_received_job_created(self): current_request = self.generate_job_request(accepted=False) new_handle = str(random.random()) self.command_handler.recv_command(GEARMAN_COMMAND_JOB_CREATED, job_handle=new_handle) self.assertEqual(current_request.job.handle, new_handle) self.assertEqual(current_request.state, JOB_CREATED) self.assertEqual(self.command_handler.handle_to_request_map[new_handle], current_request) def test_received_job_created_out_of_order(self): self.assertEqual(self.command_handler.requests_awaiting_handles, collections.deque()) # Make sure we bail cuz we have an empty queue self.assertRaises(InvalidClientState, self.command_handler.recv_command, GEARMAN_COMMAND_JOB_CREATED, job_handle=None) def test_required_state_pending(self): current_request = self.generate_job_request(submitted=False, accepted=False) new_handle = str(random.random()) invalid_states = [JOB_UNKNOWN, JOB_CREATED, JOB_COMPLETE, JOB_FAILED] for bad_state in invalid_states: current_request.state = bad_state # We only want to check the state of request... not die if we don't have any pending requests self.command_handler.requests_awaiting_handles.append(current_request) self.assertRaises(InvalidClientState, self.command_handler.recv_command, GEARMAN_COMMAND_JOB_CREATED, job_handle=new_handle) def test_required_state_queued(self): current_request = self.generate_job_request() job_handle = current_request.job.handle new_data = str(random.random()) invalid_states = [JOB_UNKNOWN, JOB_PENDING, JOB_COMPLETE, JOB_FAILED] for bad_state in invalid_states: current_request.state = bad_state # All these commands expect to be in JOB_CREATED self.assertRaises(InvalidClientState, self.command_handler.recv_command, GEARMAN_COMMAND_WORK_DATA, job_handle=job_handle, data=new_data) self.assertRaises(InvalidClientState, self.command_handler.recv_command, GEARMAN_COMMAND_WORK_WARNING, job_handle=job_handle, data=new_data) self.assertRaises(InvalidClientState, self.command_handler.recv_command, GEARMAN_COMMAND_WORK_STATUS, job_handle=job_handle, numerator=0, denominator=1) self.assertRaises(InvalidClientState, self.command_handler.recv_command, GEARMAN_COMMAND_WORK_COMPLETE, job_handle=job_handle, data=new_data) self.assertRaises(InvalidClientState, self.command_handler.recv_command, GEARMAN_COMMAND_WORK_FAIL, job_handle=job_handle) def test_in_flight_work_updates(self): current_request = self.generate_job_request() job_handle = current_request.job.handle new_data = str(random.random()) # Test WORK_DATA self.command_handler.recv_command(GEARMAN_COMMAND_WORK_DATA, job_handle=job_handle, data=new_data) self.assertEqual(current_request.data_updates.popleft(), new_data) self.assertEqual(current_request.state, JOB_CREATED) # Test WORK_WARNING self.command_handler.recv_command(GEARMAN_COMMAND_WORK_WARNING, job_handle=job_handle, data=new_data) self.assertEqual(current_request.warning_updates.popleft(), new_data) self.assertEqual(current_request.state, JOB_CREATED) # Test WORK_STATUS self.command_handler.recv_command(GEARMAN_COMMAND_WORK_STATUS, job_handle=job_handle, numerator=0, denominator=1) self.assertEqual(current_request.status_updates.popleft(), (0, 1)) self.assertEqual(current_request.state, JOB_CREATED) def test_work_complete(self): current_request = self.generate_job_request() job_handle = current_request.job.handle new_data = str(random.random()) self.command_handler.recv_command(GEARMAN_COMMAND_WORK_COMPLETE, job_handle=job_handle, data=new_data) self.assertEqual(current_request.result, new_data) self.assertEqual(current_request.state, JOB_COMPLETE) def test_work_fail(self): current_request = self.generate_job_request() job_handle = current_request.job.handle new_data = str(random.random()) self.command_handler.recv_command(GEARMAN_COMMAND_WORK_FAIL, job_handle=job_handle) self.assertEqual(current_request.state, JOB_FAILED) def test_status_request(self): current_request = self.generate_job_request() job_handle = current_request.job.handle self.assertEqual(current_request.status, {}) self.command_handler.recv_command(GEARMAN_COMMAND_STATUS_RES, job_handle=job_handle, known='1', running='1', numerator='0', denominator='1') self.assertEqual(current_request.status['handle'], job_handle) self.assertTrue(current_request.status['known']) self.assertTrue(current_request.status['running']) self.assertEqual(current_request.status['numerator'], 0) self.assertEqual(current_request.status['denominator'], 1) if __name__ == '__main__': unittest.main()