import collections
from gearman import compat
import unittest

from gearman.worker import GearmanWorker
from gearman.worker_handler import GearmanWorkerCommandHandler

from gearman.errors import ServerUnavailable, InvalidWorkerState
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

from tests._core_testing import _GearmanAbstractTest, MockGearmanConnectionManager, MockGearmanConnection

class MockGearmanWorker(MockGearmanConnectionManager, GearmanWorker):
    def __init__(self, *largs, **kwargs):
        super(MockGearmanWorker, self).__init__(*largs, **kwargs)
        self.worker_job_queues = compat.defaultdict(collections.deque)

    def on_job_execute(self, current_job):
        current_handler = self.connection_to_handler_map[current_job.connection]
        self.worker_job_queues[current_handler].append(current_job)

class _GearmanAbstractWorkerTest(_GearmanAbstractTest):
    connection_manager_class = MockGearmanWorker
    command_handler_class = GearmanWorkerCommandHandler

    def setup_command_handler(self):
        super(_GearmanAbstractWorkerTest, self).setup_command_handler()
        self.assert_sent_abilities([])
        self.assert_sent_command(GEARMAN_COMMAND_PRE_SLEEP)

    def assert_sent_abilities(self, expected_abilities):
        observed_abilities = set()

        self.assert_sent_command(GEARMAN_COMMAND_RESET_ABILITIES)
        for ability in expected_abilities:
            cmd_type, cmd_args = self.connection._outgoing_commands.popleft()

            self.assertEqual(get_command_name(cmd_type), get_command_name(GEARMAN_COMMAND_CAN_DO))
            observed_abilities.add(cmd_args['task'])

        self.assertEqual(observed_abilities, set(expected_abilities))

    def assert_sent_client_id(self, expected_client_id):
        self.assert_sent_command(GEARMAN_COMMAND_SET_CLIENT_ID, client_id=expected_client_id)

class WorkerTest(_GearmanAbstractWorkerTest):
    """Test the public worker interface"""
    def test_registering_functions(self):
        # Tests that the abilities were set on the GearmanWorker AND the GearmanWorkerCommandHandler
        # Does NOT test that commands were actually sent out as that is tested in GearmanWorkerCommandHandlerInterfaceTest.test_set_abilities
        def fake_callback_one(worker_command_handler, current_job):
            pass

        def fake_callback_two(worker_command_handler, current_job):
            pass

        # Register a single callback
        self.connection_manager.register_task('fake_callback_one', fake_callback_one)
        self.failUnless('fake_callback_one' in self.connection_manager.worker_abilities)
        self.failIf('fake_callback_two' in self.connection_manager.worker_abilities)
        self.assertEqual(self.connection_manager.worker_abilities['fake_callback_one'], fake_callback_one)
        self.assertEqual(self.command_handler._handler_abilities, ['fake_callback_one'])

        # Register another callback and make sure the command_handler sees the same functions
        self.connection_manager.register_task('fake_callback_two', fake_callback_two)
        self.failUnless('fake_callback_one' in self.connection_manager.worker_abilities)
        self.failUnless('fake_callback_two' in self.connection_manager.worker_abilities)
        self.assertEqual(self.connection_manager.worker_abilities['fake_callback_one'], fake_callback_one)
        self.assertEqual(self.connection_manager.worker_abilities['fake_callback_two'], fake_callback_two)
        self.assertEqual(self.command_handler._handler_abilities, ['fake_callback_one', 'fake_callback_two'])

        # Unregister a callback and make sure the command_handler sees the same functions
        self.connection_manager.unregister_task('fake_callback_one')
        self.failIf('fake_callback_one' in self.connection_manager.worker_abilities)
        self.failUnless('fake_callback_two' in self.connection_manager.worker_abilities)
        self.assertEqual(self.connection_manager.worker_abilities['fake_callback_two'], fake_callback_two)
        self.assertEqual(self.command_handler._handler_abilities, ['fake_callback_two'])

    def test_setting_client_id(self):
        new_client_id = 'HELLO'

        # Make sure nothing is set
        self.assertEqual(self.connection_manager.worker_client_id, None)
        self.assertEqual(self.command_handler._client_id, None)

        self.connection_manager.set_client_id(new_client_id)

        # Make sure both the client and the connection handler reflect the new state
        self.assertEqual(self.connection_manager.worker_client_id, new_client_id)
        self.assertEqual(self.command_handler._client_id, new_client_id)

    def test_establish_worker_connections(self):
        self.connection_manager.connection_list = []
        self.connection_manager.command_handlers = {}

        # Spin up a bunch of imaginary gearman connections
        good_connection = MockGearmanConnection()
        good_connection.connect()
        good_connection._fail_on_bind = False

        failed_then_retried_connection = MockGearmanConnection()
        failed_then_retried_connection._fail_on_bind = False

        failed_connection = MockGearmanConnection()
        failed_connection._fail_on_bind = True

        # Register all our connections
        self.connection_manager.connection_list = [good_connection, failed_then_retried_connection, failed_connection]

        # The only alive connections should be the ones that ultimately be connection.connected
        alive_connections = self.connection_manager.establish_worker_connections()
        self.assertTrue(good_connection in alive_connections)
        self.assertTrue(failed_then_retried_connection in alive_connections)
        self.assertFalse(failed_connection in alive_connections)

    def test_establish_worker_connections_dead(self):
        self.connection_manager.connection_list = []
        self.connection_manager.command_handlers = {}

        # We have no connections so there will never be any work to do
        self.assertRaises(ServerUnavailable, self.connection_manager.work)

        # We were started with a dead connection, make sure we bail again
        dead_connection = MockGearmanConnection()
        dead_connection._fail_on_bind = True
        dead_connection.connected = False
        self.connection_manager.connection_list = [dead_connection]

        self.assertRaises(ServerUnavailable, self.connection_manager.work)


class WorkerCommandHandlerInterfaceTest(_GearmanAbstractWorkerTest):
    """Test the public interface a GearmanWorker may need to call in order to update state on a GearmanWorkerCommandHandler"""

    def test_on_connect(self):
        expected_abilities = ['function_one', 'function_two', 'function_three']
        expected_client_id = 'my_client_id'

        self.connection.connected = False

        self.connection_manager.set_client_id(expected_client_id)
        self.connection_manager.unregister_task('__test_ability__')
        for task in expected_abilities:
            self.connection_manager.register_task(task, None)

        # We were disconnected, connect and wipe pending commands
        self.connection_manager.establish_connection(self.connection)

        # When we attempt a new connection, make sure we get a new command handler
        self.assertNotEquals(self.command_handler, self.connection_manager.connection_to_handler_map[self.connection])

        self.assert_sent_client_id(expected_client_id)
        self.assert_sent_abilities(expected_abilities)
        self.assert_sent_command(GEARMAN_COMMAND_PRE_SLEEP)
        self.assert_no_pending_commands()

    def test_set_abilities(self):
        expected_abilities = ['function_one', 'function_two', 'function_three']

        # We were disconnected, connect and wipe pending commands
        self.command_handler.set_abilities(expected_abilities)
        self.assert_sent_abilities(expected_abilities)
        self.assert_no_pending_commands()

    def test_set_client_id(self):
        expected_client_id = 'my_client_id'

        handler_initial_state = {}
        handler_initial_state['abilities'] = []
        handler_initial_state['client_id'] = None

        # We were disconnected, connect and wipe pending commands
        self.command_handler.set_client_id(expected_client_id)
        self.assert_sent_client_id(expected_client_id)
        self.assert_no_pending_commands()

    def test_send_functions(self):
        current_job = self.generate_job()

        # Test GEARMAN_COMMAND_WORK_STATUS
        self.command_handler.send_job_status(current_job, 0, 1)
        self.assert_sent_command(GEARMAN_COMMAND_WORK_STATUS, job_handle=current_job.handle, numerator='0', denominator='1')

        # Test GEARMAN_COMMAND_WORK_COMPLETE
        self.command_handler.send_job_complete(current_job, 'completion data')
        self.assert_sent_command(GEARMAN_COMMAND_WORK_COMPLETE, job_handle=current_job.handle, data='completion data')

        # Test GEARMAN_COMMAND_WORK_FAIL
        self.command_handler.send_job_failure(current_job)
        self.assert_sent_command(GEARMAN_COMMAND_WORK_FAIL, job_handle=current_job.handle)

        # Test GEARMAN_COMMAND_WORK_EXCEPTION
        self.command_handler.send_job_exception(current_job, 'exception data')
        self.assert_sent_command(GEARMAN_COMMAND_WORK_EXCEPTION, job_handle=current_job.handle, data='exception data')

        # Test GEARMAN_COMMAND_WORK_DATA
        self.command_handler.send_job_data(current_job, 'job data')
        self.assert_sent_command(GEARMAN_COMMAND_WORK_DATA, job_handle=current_job.handle, data='job data')

        # Test GEARMAN_COMMAND_WORK_WARNING
        self.command_handler.send_job_warning(current_job, 'job warning')
        self.assert_sent_command(GEARMAN_COMMAND_WORK_WARNING, job_handle=current_job.handle, data='job warning')

class WorkerCommandHandlerStateMachineTest(_GearmanAbstractWorkerTest):
    """Test multiple state transitions within a GearmanWorkerCommandHandler

    End to end tests without a server
    """
    connection_manager_class = MockGearmanWorker
    command_handler_class = GearmanWorkerCommandHandler

    def setup_connection_manager(self):
        super(WorkerCommandHandlerStateMachineTest, self).setup_connection_manager()
        self.connection_manager.register_task('__test_ability__', None)

    def setup_command_handler(self):
        super(_GearmanAbstractWorkerTest, self).setup_command_handler()
        self.assert_sent_abilities(['__test_ability__'])
        self.assert_sent_command(GEARMAN_COMMAND_PRE_SLEEP)

    def test_wakeup_work(self):
        self.move_to_state_wakeup()

        self.move_to_state_job_assign_uniq(self.generate_job_dict())

        self.move_to_state_wakeup()

        self.move_to_state_no_job()

    def test_wakeup_sleep_wakup_work(self):
        self.move_to_state_wakeup()

        self.move_to_state_no_job()

        self.move_to_state_wakeup()

        self.move_to_state_job_assign_uniq(self.generate_job_dict())

        self.move_to_state_wakeup()

        self.move_to_state_no_job()

    def test_multiple_wakeup_then_no_work(self):
        # Awaken the state machine... then give it no work
        self.move_to_state_wakeup()

        for _ in range(5):
            self.command_handler.recv_command(GEARMAN_COMMAND_NOOP)

        self.assert_job_lock(is_locked=True)

        # Pretend like the server has no work... do nothing
        # Moving to state NO_JOB will make sure there's only 1 item on the queue
        self.move_to_state_no_job()

    def test_multiple_work(self):
        self.move_to_state_wakeup()

        self.move_to_state_job_assign_uniq(self.generate_job_dict())

        self.move_to_state_wakeup()

        self.move_to_state_job_assign_uniq(self.generate_job_dict())

        self.move_to_state_wakeup()

        self.move_to_state_job_assign_uniq(self.generate_job_dict())

        self.move_to_state_wakeup()

        # After this job completes, we're going to greedily ask for more jobs
        self.move_to_state_no_job()

    def test_worker_already_locked(self):
        other_connection = MockGearmanConnection()
        self.connection_manager.connection_list.append(other_connection)
        self.connection_manager.establish_connection(other_connection)

        other_handler = self.connection_manager.connection_to_handler_map[other_connection]
        other_handler.recv_command(GEARMAN_COMMAND_NOOP)

        # Make sure other handler has a lock
        self.assertEqual(self.connection_manager.command_handler_holding_job_lock, other_handler)

        # Make sure OUR handler has nothing incoming
        self.assert_no_pending_commands()

        # Make sure we try to grab a job but fail...so go back to sleep
        self.command_handler.recv_command(GEARMAN_COMMAND_NOOP)
        self.assert_sent_command(GEARMAN_COMMAND_PRE_SLEEP)

        # Make sure other handler still has lock
        self.assertEqual(self.connection_manager.command_handler_holding_job_lock, other_handler)

        # Make the other handler release its lock
        other_handler.recv_command(GEARMAN_COMMAND_NO_JOB)

        # Ensure that the lock has been freed
        self.assert_job_lock(is_locked=False)

        # Try to do work after we have our lock released
        self.move_to_state_wakeup()

        self.move_to_state_job_assign_uniq(self.generate_job_dict())

        self.move_to_state_wakeup()

        self.move_to_state_no_job()

    def move_to_state_wakeup(self):
        self.assert_no_pending_commands()
        self.assert_job_lock(is_locked=False)

        self.command_handler.recv_command(GEARMAN_COMMAND_NOOP)

    def move_to_state_no_job(self):
        """Move us to the NO_JOB state...

        1) We should've most recently sent only a single GEARMAN_COMMAND_GRAB_JOB_UNIQ
        2) We should be awaiting job assignment
        3) Once we receive a NO_JOB, we should say we're going back to sleep"""
        self.assert_awaiting_job()

        self.command_handler.recv_command(GEARMAN_COMMAND_NO_JOB)

        # We should be asleep... which means no pending jobs and we're not awaiting job assignment
        self.assert_sent_command(GEARMAN_COMMAND_PRE_SLEEP)
        self.assert_no_pending_commands()
        self.assert_job_lock(is_locked=False)

    def move_to_state_job_assign_uniq(self, fake_job):
        """Move us to the JOB_ASSIGN_UNIQ state...

        1) We should've most recently sent only a single GEARMAN_COMMAND_GRAB_JOB_UNIQ
        2) We should be awaiting job assignment
        3) The job we receive should be the one we expected"""
        self.assert_awaiting_job()

        ### NOTE: This recv_command does NOT send out a GEARMAN_COMMAND_JOB_COMPLETE or GEARMAN_COMMAND_JOB_FAIL
        ###           as we're using a MockGearmanConnectionManager with a method that only queues the job
        self.command_handler.recv_command(GEARMAN_COMMAND_JOB_ASSIGN_UNIQ, **fake_job)

        current_job = self.connection_manager.worker_job_queues[self.command_handler].popleft()
        self.assertEqual(current_job.handle, fake_job['job_handle'])
        self.assertEqual(current_job.task, fake_job['task'])
        self.assertEqual(current_job.unique, fake_job['unique'])
        self.assertEqual(current_job.data, fake_job['data'])

        # At the end of recv_command(GEARMAN_COMMAND_JOB_ASSIGN_UNIQ)
        self.assert_job_lock(is_locked=False)
        self.assert_sent_command(GEARMAN_COMMAND_PRE_SLEEP)

    def assert_awaiting_job(self):
        self.assert_sent_command(GEARMAN_COMMAND_GRAB_JOB_UNIQ)
        self.assert_no_pending_commands()

    def assert_job_lock(self, is_locked):
        expected_value = (is_locked and self.command_handler) or None
        self.assertEqual(self.connection_manager.command_handler_holding_job_lock, expected_value)

if __name__ == '__main__':
    unittest.main()