def test_create_datetime(self):

        comparator = datetime.datetime.now(tz=dateutils.utc_tz())
        result = dateutils.now_utc_datetime_with_tzinfo()
        self.assertTrue(hasattr(result, 'tzinfo'))
        self.assertEquals(result.tzinfo, dateutils.utc_tz())
        self.assertTrue(result >= comparator)

def convert_schedule(save_func, call):
    """
    Converts one scheduled call from the old schema to the new

    :param save_func:   a function that takes one parameter, a dictionary that
                        represents the scheduled call in its new schema. This
                        function should save the call to the database.
    :type  save_func:   function
    :param call:        dictionary representing the scheduled call in its old
                        schema
    :type  call:        dict
    """
    call.pop('call_exit_states', None)
    call['total_run_count'] = call.pop('call_count')

    call['iso_schedule'] = call['schedule']
    interval, start_time, occurrences = dateutils.parse_iso8601_interval(call['schedule'])
    # this should be a pickled instance of celery.schedules.schedule
    call['schedule'] = pickle.dumps(schedule(interval))

    call_request = call.pop('serialized_call_request')
    # we are no longer storing these pickled.
    # these are cast to a string because python 2.6 sometimes fails to
    # deserialize json from unicode.
    call['args'] = pickle.loads(str(call_request['args']))
    call['kwargs'] = pickle.loads(str(call_request['kwargs']))
    # keeping this pickled because we don't really know how to use it yet
    call['principal'] = call_request['principal']
    # this always get calculated on-the-fly now
    call.pop('next_run', None)
    first_run = call['first_run'].replace(tzinfo=dateutils.utc_tz())
    call['first_run'] = dateutils.format_iso8601_datetime(first_run)
    last_run = call.pop('last_run')
    if last_run:
        last_run_at = last_run.replace(tzinfo=dateutils.utc_tz())
        call['last_run_at'] = dateutils.format_iso8601_datetime(last_run_at)
    else:
        call['last_run_at'] = None
    call['task'] = NAMES_TO_TASKS[call_request['callable_name']]

    # this is a new field that is used to determine when the scheduler needs to
    # re-read the collection of schedules.
    call['last_updated'] = time.time()

    # determine if this is a consumer-related schedule, which we can only identify
    # by the consumer resource tag. If it is, save that tag value in the new
    # "resource" field, which is the new way that we will identify the
    # relationship between a schedule and some other object. This is not
    # necessary for repos, because we have a better method above for identifying
    # them (move_scheduled_syncs).
    tags = call_request.get('tags', [])
    for tag in tags:
        if tag.startswith('pulp:consumer:'):
            call['resource'] = tag
            break

    save_func(call)

 def _run(self):
     """
     Run the call in the call request.
     Generally the target of a new thread.
     """
     # used for calling _run directly during testing
     principal_manager = managers_factory.principal_manager()
     principal_manager.set_principal(self.call_request.principal)
     # usually set in the wrapper, unless called directly
     if self.call_report.state in dispatch_constants.CALL_READY_STATES:
         self.call_report.state = dispatch_constants.CALL_RUNNING_STATE
     self.call_report.start_time = datetime.datetime.now(dateutils.utc_tz())
     dispatch_context.CONTEXT.set_task_attributes(self)
     call = self.call_request.call
     args = copy.copy(self.call_request.args)
     kwargs = copy.copy(self.call_request.kwargs)
     try:
         result = call(*args, **kwargs)
     except:
         # NOTE: this is making an assumption here that the call failed to
         # execute, if this isn't the case, or it got far enough, we may be
         # faced with _succeeded or _failed being called again
         e, tb = sys.exc_info()[1:]
         _LOG.exception(e)
         # too bad 2.4 doesn't support try/except/finally blocks
         principal_manager.clear_principal()
         dispatch_context.CONTEXT.clear_task_attributes()
         return self._failed(e, tb)
     principal_manager.clear_principal()
     dispatch_context.CONTEXT.clear_task_attributes()
     return result

    def test_last_publish(self):
        """
        Tests retrieving the last publish time in both the unpublish and previously published cases.
        """

        class GMT5(datetime.tzinfo):
            def utcoffset(self, dt):
                return datetime.timedelta(hours=5, minutes=30)

            def tzname(self, dt):
                return "GMT +5"

            def dst(self, dt):
                return datetime.timedelta(0)

        # Test - Unpublished
        unpublished = self.conduit.last_publish()
        self.assertTrue(unpublished is None)

        # Setup - Previous publish
        last_publish = datetime.datetime(2015, 4, 29, 20, 23, 56, 0, tzinfo=GMT5())
        repo_dist = model.Distributor.objects.get_or_404(repo_id='repo-1')
        repo_dist['last_publish'] = last_publish
        repo_dist.save()

        # Test - Last publish
        found = self.conduit.last_publish()
        self.assertTrue(isinstance(found, datetime.datetime))  # check returned format

        self.assertEqual(found.tzinfo, dateutils.utc_tz())
        self.assertEqual(repo_dist['last_publish'], found)

 def status(cls, uuids=[]):
     """
     Get the agent heartbeat status.
     @param uuids: An (optional) list of uuids to query.
     @return: A tuple (status,last-heartbeat)
     """
     cls.__lock()
     try:
         now = dt.now(dateutils.utc_tz())
         if not uuids:
             uuids = cls.__status.keys()
         d = {}
         for uuid in uuids:
             last = cls.__status.get(uuid)
             if last:
                 status = ( last[1] > now )
                 heartbeat = last[0].isoformat()
                 any = last[2]
             else:
                 status = False
                 heartbeat = None
                 any = {}
             d[uuid] = (status, heartbeat, any)
         return d
     finally:
         cls.__unlock()

    def set_task_failed(task_id, traceback=None, timestamp=None):
        """
        Update a task's state to reflect that it has succeeded.
        :param task_id: The identity of the task to be updated.
        :type  task_id: basestring
        :ivar traceback: A string representation of the traceback resulting from the task execution.
        :type traceback: basestring
        :param timestamp: The (optional) ISO-8601 finished timestamp (UTC).
        :type timestamp: str
        """
        collection = TaskStatus.get_collection()

        if not timestamp:
            now = datetime.now(dateutils.utc_tz())
            finished = dateutils.format_iso8601_datetime(now)
        else:
            finished = timestamp

        update = {
            '$set': {
                'finish_time': finished,
                'state': constants.CALL_ERROR_STATE,
                'traceback': traceback
            }
        }

        collection.update({'task_id': task_id}, update, safe=True)

 def __call__(self, *args, **kwargs):
     """
     This overrides CeleryTask's __call__() method. We use this method
     for task state tracking of Pulp tasks.
     """
     # Check task status and skip running the task if task state is 'canceled'.
     try:
         task_status = TaskStatus.objects.get(task_id=self.request.id)
     except DoesNotExist:
         task_status = None
     if task_status and task_status['state'] == constants.CALL_CANCELED_STATE:
         _logger.debug("Task cancel received for task-id : [%s]" % self.request.id)
         return
     # Update start_time and set the task state to 'running' for asynchronous tasks.
     # Skip updating status for eagerly executed tasks, since we don't want to track
     # synchronous tasks in our database.
     if not self.request.called_directly:
         now = datetime.now(dateutils.utc_tz())
         start_time = dateutils.format_iso8601_datetime(now)
         # Using 'upsert' to avoid a possible race condition described in the apply_async method
         # above.
         TaskStatus.objects(task_id=self.request.id).update_one(
             set__state=constants.CALL_RUNNING_STATE, set__start_time=start_time, upsert=True)
     # Run the actual task
     _logger.debug("Running task : [%s]" % self.request.id)
     return super(Task, self).__call__(*args, **kwargs)

    def _run(self):
        """
        Run the call in the call request.
        Generally the target of a new thread.
        """
        # used for calling _run directly during testing
        principal_manager = managers_factory.principal_manager()
        principal_manager.set_principal(self.call_request.principal)

        # generally set in the wrapper, but not when called directly
        if self.call_report.state in dispatch_constants.CALL_READY_STATES:
            self.call_report.state = dispatch_constants.CALL_RUNNING_STATE

        self.call_report.start_time = datetime.datetime.now(dateutils.utc_tz())

        dispatch_context.CONTEXT.set_task_attributes(self)

        call = self.call_request.call
        args = copy.copy(self.call_request.args)
        kwargs = copy.copy(self.call_request.kwargs)

        try:
            result = call(*args, **kwargs)

        except:
            e, tb = sys.exc_info()[1:]
            _LOG.exception(e)
            return self._failed(e, tb)

        else:
            return self._succeeded(result)

        finally:
            principal_manager.clear_principal()
            dispatch_context.CONTEXT.clear_task_attributes()

    def test_updated_scheduled_next_run(self):
        call_request = CallRequest(itinerary_call)
        interval = datetime.timedelta(minutes=2)
        now = datetime.datetime.now(tz=dateutils.utc_tz())
        old_schedule = dateutils.format_iso8601_interval(interval, now)

        scheduled_id = self.scheduler.add(call_request, old_schedule)

        self.assertNotEqual(scheduled_id, None)

        scheduled_call = self.scheduled_call_collection.find_one({'_id': ObjectId(scheduled_id)})

        self.assertNotEqual(scheduled_call, None)

        old_interval, start_time = dateutils.parse_iso8601_interval(old_schedule)[:2]
        start_time = dateutils.to_naive_utc_datetime(start_time)

        self.assertEqual(scheduled_call['last_run'], None)
        self.assertEqual(scheduled_call['first_run'], start_time + old_interval)
        self.assertEqual(scheduled_call['next_run'], start_time + old_interval)

        interval = datetime.timedelta(minutes=1)
        new_schedule = dateutils.format_iso8601_interval(interval, now)

        self.scheduler.update(scheduled_id, schedule=new_schedule)
        updated_scheduled_call = self.scheduled_call_collection.find_one({'_id': ObjectId(scheduled_id)})

        new_interval = dateutils.parse_iso8601_interval(new_schedule)[0]

        self.assertEqual(updated_scheduled_call['last_run'], None)
        self.assertEqual(updated_scheduled_call['first_run'], start_time + old_interval)
        self.assertEqual(updated_scheduled_call['next_run'], start_time + new_interval)

    def test_task_status_update_fires_notification(self, mock_send):
        """
        Test that update_one() also fires a notification.
        """
        task_id = self.get_random_uuid()
        worker_name = 'special_worker_name'
        tags = ['test-tag1', 'test-tag2']
        state = 'waiting'
        ts = TaskStatus(task_id, worker_name, tags, state)
        ts.save()
        # ensure event was fired for save()
        mock_send.assert_called_once_with(ts, routing_key="tasks.%s" % task_id)
        now = datetime.now(dateutils.utc_tz())
        start_time = dateutils.format_iso8601_datetime(now)
        delta = {'start_time': start_time,
                 'state': 'running',
                 'progress_report': {'report-id': 'my-progress'}}

        self.assertEquals(len(mock_send.call_args_list), 1)
        TaskStatus.objects(task_id=task_id).update_one(
            set__start_time=delta['start_time'], set__state=delta['state'],
            set__progress_report=delta['progress_report'])

        # ensure event was fired for update_one()
        self.assertEquals(len(mock_send.call_args_list), 2)
        mock_send.assert_called_with(ts, routing_key="tasks.%s" % task_id)

    def on_failure(self, exc, task_id, args, kwargs, einfo):
        """
        This overrides the error handler run by the worker when the task fails.
        It updates state, finish_time and traceback of the relevant task status
        for asynchronous tasks. Skip updating status for synchronous tasks.

        :param exc:     The exception raised by the task.
        :param task_id: Unique id of the failed task.
        :param args:    Original arguments for the executed task.
        :param kwargs:  Original keyword arguments for the executed task.
        :param einfo:   celery's ExceptionInfo instance, containing serialized traceback.
        """
        if isinstance(exc, PulpCodedException):
            _logger.info(_('Task failed : [%(task_id)s] : %(msg)s') %
                         {'task_id': task_id, 'msg': str(exc)})
            _logger.debug(traceback.format_exc())
        else:
            _logger.info(_('Task failed : [%s]') % task_id)
            # celery will log the traceback
        if not self.request.called_directly:
            now = datetime.now(dateutils.utc_tz())
            finish_time = dateutils.format_iso8601_datetime(now)
            task_status = TaskStatus.objects.get(task_id=task_id)
            task_status['state'] = constants.CALL_ERROR_STATE
            task_status['finish_time'] = finish_time
            task_status['traceback'] = einfo.traceback
            if not isinstance(exc, PulpException):
                exc = PulpException(str(exc))
            task_status['error'] = exc.to_dict()

            task_status.save()
            common_utils.delete_working_directory()

    def failed(self, reply):
        """
        Notification (reply) indicating an RMI failed.
        This information used to update the task status.
        :param reply: A failure reply object.
        :type reply: gofer.rmi.async.Failed
        """
        _logger.info(_('Task RMI (failed): %(r)s'), {'r': reply})

        call_context = dict(reply.data)
        action = call_context.get('action')
        task_id = call_context['task_id']
        traceback = reply.xstate['trace']
        finished = reply.timestamp
        if not finished:
            now = datetime.now(dateutils.utc_tz())
            finished = dateutils.format_iso8601_datetime(now)

        TaskStatus.objects(task_id=task_id).update_one(set__finish_time=finished,
                                                       set__state=constants.CALL_ERROR_STATE,
                                                       set__traceback=traceback)

        if action == 'bind':
            ReplyHandler._bind_failed(task_id, call_context)
            return
        if action == 'unbind':
            ReplyHandler._unbind_failed(task_id, call_context)
            return

 def test_ensure_tz_not_specified(self):
     """
     Make sure that a date without a timezone is given one.
     """
     dt = datetime.datetime.utcnow()
     new_date = Repository._ensure_tz_specified(dt)
     self.assertEquals(new_date.tzinfo, dateutils.utc_tz())

 def test_tz_not_specified(self):
     """
     Test that if a tz is not specified, it is added.
     """
     dt = datetime.datetime.utcnow()
     new_date = dateutils.ensure_tz(dt)
     self.assertEquals(new_date.tzinfo, dateutils.utc_tz())

 def test_tz_specified(self):
     """
     Ensure that if the tz is already specified, it is used.
     """
     dt = datetime.datetime.now(dateutils.local_tz())
     new_date = dateutils.ensure_tz(dt)
     self.assertEquals(new_date.tzinfo, dateutils.utc_tz())

    def test_now(self):
        call = ScheduledCall('PT1H', 'pulp.tasks.dosomething')

        now = datetime.utcnow().replace(tzinfo=dateutils.utc_tz())
        next_run = dateutils.parse_iso8601_datetime(call.calculate_next_run())

        self.assertTrue(next_run - now < timedelta(seconds=1))

    def set_task_succeeded(task_id, result=None, timestamp=None):
        """
        Update a task's state to reflect that it has succeeded.
        :param task_id: The identity of the task to be updated.
        :type  task_id: basestring
        :param result: The optional value returned by the task execution.
        :type result: anything
        :param timestamp: The (optional) ISO-8601 finished timestamp (UTC).
        :type timestamp: str
        """
        collection = TaskStatus.get_collection()

        if not timestamp:
            now = datetime.now(dateutils.utc_tz())
            finished = dateutils.format_iso8601_datetime(now)
        else:
            finished = timestamp

        update = {
            '$set': {
                'finish_time': finished,
                'state': constants.CALL_FINISHED_STATE,
                'result': result
            }
        }

        collection.update({'task_id': task_id}, update, safe=True)

    def test_task_status_update(self):
        """
        Tests the successful operation of task status update.
        """
        task_id = self.get_random_uuid()
        worker_name = 'special_worker_name'
        tags = ['test-tag1', 'test-tag2']
        state = 'waiting'
        TaskStatus(task_id, worker_name, tags, state).save()
        now = datetime.now(dateutils.utc_tz())
        start_time = dateutils.format_iso8601_datetime(now)
        delta = {'start_time': start_time,
                 'state': 'running',
                 'progress_report': {'report-id': 'my-progress'}}

        TaskStatus.objects(task_id=task_id).update_one(
            set__start_time=delta['start_time'], set__state=delta['state'],
            set__progress_report=delta['progress_report'])

        task_status = TaskStatus.objects(task_id=task_id).first()
        self.assertEqual(task_status['start_time'], delta['start_time'])
        # Make sure that parse_iso8601_datetime is able to parse the start_time without errors
        dateutils.parse_iso8601_datetime(task_status['start_time'])
        self.assertEqual(task_status['state'], delta['state'])
        self.assertEqual(task_status['progress_report'], delta['progress_report'])
        self.assertEqual(task_status['worker_name'], worker_name)

    def test_update_task_status(self):
        """
        Tests the successful operation of update_task_status().
        """
        task_id = self.get_random_uuid()
        queue = 'special_queue'
        tags = ['test-tag1', 'test-tag2']
        state = 'waiting'
        TaskStatusManager.create_task_status(task_id, queue, tags, state)
        now = datetime.now(dateutils.utc_tz())
        start_time = dateutils.format_iso8601_datetime(now)
        delta = {'start_time': start_time,
                 'state': 'running',
                 'disregard': 'ignored',
                 'progress_report': {'report-id': 'my-progress'}}

        updated = TaskStatusManager.update_task_status(task_id, delta)

        task_status = TaskStatusManager.find_by_task_id(task_id)
        self.assertEqual(task_status['start_time'], delta['start_time'])
        # Make sure that parse_iso8601_datetime is able to parse the start_time without errors
        dateutils.parse_iso8601_datetime(task_status['start_time'])
        self.assertEqual(task_status['state'], delta['state'])
        self.assertEqual(task_status['progress_report'], delta['progress_report'])
        self.assertEqual(task_status['queue'], queue)
        self.assertEqual(updated['start_time'], delta['start_time'])
        self.assertEqual(updated['state'], delta['state'])
        self.assertEqual(updated['progress_report'], delta['progress_report'])
        self.assertTrue('disregard' not in updated)
        self.assertTrue('disregard' not in task_status)

    def succeeded(self, reply):
        """
        Notification (reply) indicating an RMI succeeded.
        This information is relayed to the task coordinator.
        :param reply: A successful reply object.
        :type reply: gofer.rmi.async.Succeeded
        """
        _logger.info(_('Task RMI (succeeded): %(r)s'), {'r': reply})

        call_context = dict(reply.data)
        action = call_context.get('action')
        task_id = call_context['task_id']
        result = dict(reply.retval)
        finished = reply.timestamp
        if not finished:
            now = datetime.now(dateutils.utc_tz())
            finished = dateutils.format_iso8601_datetime(now)

        TaskStatus.objects(task_id=task_id).update_one(set__finish_time=finished,
                                                       set__state=constants.CALL_FINISHED_STATE,
                                                       set__result=result)
        if action == 'bind':
            if result['succeeded']:
                ReplyHandler._bind_succeeded(task_id, call_context)
            else:
                ReplyHandler._bind_failed(task_id, call_context)
            return
        if action == 'unbind':
            if result['succeeded']:
                ReplyHandler._unbind_succeeded(call_context)
            else:
                ReplyHandler._unbind_failed(task_id, call_context)
            return