def test_cancel_with_provided_exception(self):
     message = 'my cancel message'
     transfer_coordinator = TransferCoordinator()
     self.coordinator_controller.add_transfer_coordinator(
         transfer_coordinator)
     self.coordinator_controller.cancel(message, exc_type=FatalError)
     transfer_coordinator.announce_done()
     with self.assertRaisesRegexp(FatalError, message):
         transfer_coordinator.result()
 def test_cancel_with_message(self):
     message = 'my cancel message'
     transfer_coordinator = TransferCoordinator()
     self.coordinator_controller.add_transfer_coordinator(
         transfer_coordinator)
     self.coordinator_controller.cancel(message)
     transfer_coordinator.announce_done()
     with self.assertRaisesRegexp(CancelledError, message):
         transfer_coordinator.result()
class TestTask(unittest.TestCase):
    def setUp(self):
        self.transfer_id = 1
        self.transfer_coordinator = TransferCoordinator(
            transfer_id=self.transfer_id)

    def test_repr(self):
        main_kwargs = {'bucket': 'mybucket', 'param_to_not_include': 'foo'}
        task = ReturnKwargsTask(self.transfer_coordinator,
                                main_kwargs=main_kwargs)
        # The repr should not include the other parameter because it is not
        # a desired parameter to include.
        self.assertEqual(
            repr(task),
            'ReturnKwargsTask(transfer_id=%s, %s)' % (self.transfer_id, {
                'bucket': 'mybucket'
            }))

    def test_transfer_id(self):
        task = SuccessTask(self.transfer_coordinator)
        # Make sure that the id is the one provided to the id associated
        # to the transfer coordinator.
        self.assertEqual(task.transfer_id, self.transfer_id)

    def test_context_status_transitioning_success(self):
        # The status should be set to running.
        self.transfer_coordinator.set_status_to_running()
        self.assertEqual(self.transfer_coordinator.status, 'running')

        # If a task is called, the status still should be running.
        SuccessTask(self.transfer_coordinator)()
        self.assertEqual(self.transfer_coordinator.status, 'running')

        # Once the final task is called, the status should be set to success.
        SuccessTask(self.transfer_coordinator, is_final=True)()
        self.assertEqual(self.transfer_coordinator.status, 'success')

    def test_context_status_transitioning_failed(self):
        self.transfer_coordinator.set_status_to_running()

        SuccessTask(self.transfer_coordinator)()
        self.assertEqual(self.transfer_coordinator.status, 'running')

        # A failure task should result in the failed status
        FailureTask(self.transfer_coordinator)()
        self.assertEqual(self.transfer_coordinator.status, 'failed')

        # Even if the final task comes in and succeeds, it should stay failed.
        SuccessTask(self.transfer_coordinator, is_final=True)()
        self.assertEqual(self.transfer_coordinator.status, 'failed')

    def test_result_setting_for_success(self):
        override_return = 'foo'
        SuccessTask(self.transfer_coordinator)()
        SuccessTask(self.transfer_coordinator,
                    main_kwargs={'return_value': override_return},
                    is_final=True)()

        # The return value for the transfer future should be of the final
        # task.
        self.assertEqual(self.transfer_coordinator.result(), override_return)

    def test_result_setting_for_error(self):
        FailureTask(self.transfer_coordinator)()

        # If another failure comes in, the result should still throw the
        # original exception when result() is eventually called.
        FailureTask(self.transfer_coordinator,
                    main_kwargs={'exception': Exception})()

        # Even if a success task comes along, the result of the future
        # should be the original exception
        SuccessTask(self.transfer_coordinator, is_final=True)()
        with self.assertRaises(TaskFailureException):
            self.transfer_coordinator.result()

    def test_done_callbacks_success(self):
        callback_results = []
        SuccessTask(self.transfer_coordinator,
                    done_callbacks=[
                        partial(callback_results.append, 'first'),
                        partial(callback_results.append, 'second')
                    ])()
        # For successful tasks, the done callbacks should get called.
        self.assertEqual(callback_results, ['first', 'second'])

    def test_done_callbacks_failure(self):
        callback_results = []
        FailureTask(self.transfer_coordinator,
                    done_callbacks=[
                        partial(callback_results.append, 'first'),
                        partial(callback_results.append, 'second')
                    ])()
        # For even failed tasks, the done callbacks should get called.
        self.assertEqual(callback_results, ['first', 'second'])

        # Callbacks should continue to be called even after a related failure
        SuccessTask(self.transfer_coordinator,
                    done_callbacks=[
                        partial(callback_results.append, 'third'),
                        partial(callback_results.append, 'fourth')
                    ])()
        self.assertEqual(callback_results,
                         ['first', 'second', 'third', 'fourth'])

    def test_failure_cleanups_on_failure(self):
        callback_results = []
        self.transfer_coordinator.add_failure_cleanup(callback_results.append,
                                                      'first')
        self.transfer_coordinator.add_failure_cleanup(callback_results.append,
                                                      'second')
        FailureTask(self.transfer_coordinator)()
        # The failure callbacks should have not been called yet because it
        # is not the last task
        self.assertEqual(callback_results, [])

        # Now the failure callbacks should get called.
        SuccessTask(self.transfer_coordinator, is_final=True)()
        self.assertEqual(callback_results, ['first', 'second'])

    def test_no_failure_cleanups_on_success(self):
        callback_results = []
        self.transfer_coordinator.add_failure_cleanup(callback_results.append,
                                                      'first')
        self.transfer_coordinator.add_failure_cleanup(callback_results.append,
                                                      'second')
        SuccessTask(self.transfer_coordinator, is_final=True)()
        # The failure cleanups should not have been called because no task
        # failed for the transfer context.
        self.assertEqual(callback_results, [])

    def test_passing_main_kwargs(self):
        main_kwargs = {'foo': 'bar', 'baz': 'biz'}
        ReturnKwargsTask(self.transfer_coordinator,
                         main_kwargs=main_kwargs,
                         is_final=True)()
        # The kwargs should have been passed to the main()
        self.assertEqual(self.transfer_coordinator.result(), main_kwargs)

    def test_passing_pending_kwargs_single_futures(self):
        pending_kwargs = {}
        ref_main_kwargs = {'foo': 'bar', 'baz': 'biz'}

        # Pass some tasks to an executor
        with futures.ThreadPoolExecutor(1) as executor:
            pending_kwargs['foo'] = executor.submit(
                SuccessTask(
                    self.transfer_coordinator,
                    main_kwargs={'return_value': ref_main_kwargs['foo']}))
            pending_kwargs['baz'] = executor.submit(
                SuccessTask(
                    self.transfer_coordinator,
                    main_kwargs={'return_value': ref_main_kwargs['baz']}))

        # Create a task that depends on the tasks passed to the executor
        ReturnKwargsTask(self.transfer_coordinator,
                         pending_main_kwargs=pending_kwargs,
                         is_final=True)()
        # The result should have the pending keyword arg values flushed
        # out.
        self.assertEqual(self.transfer_coordinator.result(), ref_main_kwargs)

    def test_passing_pending_kwargs_list_of_futures(self):
        pending_kwargs = {}
        ref_main_kwargs = {'foo': ['first', 'second']}

        # Pass some tasks to an executor
        with futures.ThreadPoolExecutor(1) as executor:
            first_future = executor.submit(
                SuccessTask(
                    self.transfer_coordinator,
                    main_kwargs={'return_value': ref_main_kwargs['foo'][0]}))
            second_future = executor.submit(
                SuccessTask(
                    self.transfer_coordinator,
                    main_kwargs={'return_value': ref_main_kwargs['foo'][1]}))
            # Make the pending keyword arg value a list
            pending_kwargs['foo'] = [first_future, second_future]

        # Create a task that depends on the tasks passed to the executor
        ReturnKwargsTask(self.transfer_coordinator,
                         pending_main_kwargs=pending_kwargs,
                         is_final=True)()
        # The result should have the pending keyword arg values flushed
        # out in the expected order.
        self.assertEqual(self.transfer_coordinator.result(), ref_main_kwargs)

    def test_passing_pending_and_non_pending_kwargs(self):
        main_kwargs = {'nonpending_value': 'foo'}
        pending_kwargs = {}
        ref_main_kwargs = {
            'nonpending_value': 'foo',
            'pending_value': 'bar',
            'pending_list': ['first', 'second']
        }

        # Create the pending tasks
        with futures.ThreadPoolExecutor(1) as executor:
            pending_kwargs['pending_value'] = executor.submit(
                SuccessTask(self.transfer_coordinator,
                            main_kwargs={
                                'return_value':
                                ref_main_kwargs['pending_value']
                            }))

            first_future = executor.submit(
                SuccessTask(self.transfer_coordinator,
                            main_kwargs={
                                'return_value':
                                ref_main_kwargs['pending_list'][0]
                            }))
            second_future = executor.submit(
                SuccessTask(self.transfer_coordinator,
                            main_kwargs={
                                'return_value':
                                ref_main_kwargs['pending_list'][1]
                            }))
            # Make the pending keyword arg value a list
            pending_kwargs['pending_list'] = [first_future, second_future]

        # Create a task that depends on the tasks passed to the executor
        # and just regular nonpending kwargs.
        ReturnKwargsTask(self.transfer_coordinator,
                         main_kwargs=main_kwargs,
                         pending_main_kwargs=pending_kwargs,
                         is_final=True)()
        # The result should have all of the kwargs (both pending and
        # nonpending)
        self.assertEqual(self.transfer_coordinator.result(), ref_main_kwargs)

    def test_single_failed_pending_future(self):
        pending_kwargs = {}

        # Pass some tasks to an executor. Make one successful and the other
        # a failure.
        with futures.ThreadPoolExecutor(1) as executor:
            pending_kwargs['foo'] = executor.submit(
                SuccessTask(self.transfer_coordinator,
                            main_kwargs={'return_value': 'bar'}))
            pending_kwargs['baz'] = executor.submit(
                FailureTask(self.transfer_coordinator))

        # Create a task that depends on the tasks passed to the executor
        ReturnKwargsTask(self.transfer_coordinator,
                         pending_main_kwargs=pending_kwargs,
                         is_final=True)()
        # The end result should raise the exception from the initial
        # pending future value
        with self.assertRaises(TaskFailureException):
            self.transfer_coordinator.result()

    def test_single_failed_pending_future_in_list(self):
        pending_kwargs = {}

        # Pass some tasks to an executor. Make one successful and the other
        # a failure.
        with futures.ThreadPoolExecutor(1) as executor:
            first_future = executor.submit(
                SuccessTask(self.transfer_coordinator,
                            main_kwargs={'return_value': 'bar'}))
            second_future = executor.submit(
                FailureTask(self.transfer_coordinator))

            pending_kwargs['pending_list'] = [first_future, second_future]

        # Create a task that depends on the tasks passed to the executor
        ReturnKwargsTask(self.transfer_coordinator,
                         pending_main_kwargs=pending_kwargs,
                         is_final=True)()
        # The end result should raise the exception from the initial
        # pending future value in the list
        with self.assertRaises(TaskFailureException):
            self.transfer_coordinator.result()