def test_update(self):
        saved = TraceDBTest._create_save_trace(
            trace_tag='test_trace',
            action_executions=[],
            rules=[],
            trigger_instances=[])
        retrieved = Trace.query(trace_tag=saved.trace_tag)
        self.assertEquals(len(retrieved), 1, 'Should have 1 trace.')
        self.assertEquals(retrieved[0].id, saved.id, 'Incorrect trace retrieved.')

        no_action_executions = 4
        no_rules = 4
        no_trigger_instances = 5
        saved = TraceDBTest._create_save_trace(
            trace_tag='test_trace',
            id_=retrieved[0].id,
            action_executions=[str(bson.ObjectId()) for _ in range(no_action_executions)],
            rules=[str(bson.ObjectId()) for _ in range(no_rules)],
            trigger_instances=[str(bson.ObjectId()) for _ in range(no_trigger_instances)])
        retrieved = Trace.query(trace_tag=saved.trace_tag)

        self.assertEquals(len(retrieved), 1, 'Should have 1 trace.')
        self.assertEquals(retrieved[0].id, saved.id, 'Incorrect trace retrieved.')
        # validate update
        self.assertEquals(len(retrieved[0].action_executions), no_action_executions,
                          'Failed to update action_executions.')
        self.assertEquals(len(retrieved[0].rules), no_rules, 'Failed to update rules.')
        self.assertEquals(len(retrieved[0].trigger_instances), no_trigger_instances,
                          'Failed to update trigger_instances.')

    def test_add_or_update_given_trace_context_new(self):
        trace_context = {'trace_tag': 'awesome_test_trace'}
        action_execution_id = 'action_execution_1'
        rule_id = 'rule_1'
        trigger_instance_id = 'trigger_instance_1'

        pre_add_or_update_traces = len(Trace.get_all())
        trace_db = trace_service.add_or_update_given_trace_context(
            trace_context,
            action_executions=[action_execution_id],
            rules=[rule_id],
            trigger_instances=[trigger_instance_id])
        post_add_or_update_traces = len(Trace.get_all())

        self.assertTrue(post_add_or_update_traces > pre_add_or_update_traces,
                        'Expected new Trace to be created.')

        retrieved_trace_db = Trace.get_by_id(trace_db.id)
        self.assertEqual(len(retrieved_trace_db.action_executions), 1,
                         'Expected updated action_executions.')
        self.assertEqual(retrieved_trace_db.action_executions[0].object_id, action_execution_id,
                         'Expected updated action_executions.')

        self.assertEqual(len(retrieved_trace_db.rules), 1, 'Expected updated rules.')
        self.assertEqual(retrieved_trace_db.rules[0].object_id, rule_id, 'Expected updated rules.')

        self.assertEqual(len(retrieved_trace_db.trigger_instances), 1,
                         'Expected updated trigger_instances.')
        self.assertEqual(retrieved_trace_db.trigger_instances[0].object_id, trigger_instance_id,
                         'Expected updated trigger_instances.')

        Trace.delete(retrieved_trace_db)

    def test_add_or_update_given_trace_context(self):
        trace_context = {'id_': str(self.trace_empty.id)}
        action_execution_id = 'action_execution_1'
        rule_id = 'rule_1'
        trigger_instance_id = 'trigger_instance_1'
        trace_service.add_or_update_given_trace_context(
            trace_context,
            action_executions=[action_execution_id],
            rules=[rule_id],
            trigger_instances=[trigger_instance_id])

        retrieved_trace_db = Trace.get_by_id(self.trace_empty.id)
        self.assertEqual(len(retrieved_trace_db.action_executions), 1,
                         'Expected updated action_executions.')
        self.assertEqual(retrieved_trace_db.action_executions[0].object_id, action_execution_id,
                         'Expected updated action_executions.')

        self.assertEqual(len(retrieved_trace_db.rules), 1, 'Expected updated rules.')
        self.assertEqual(retrieved_trace_db.rules[0].object_id, rule_id, 'Expected updated rules.')

        self.assertEqual(len(retrieved_trace_db.trigger_instances), 1,
                         'Expected updated trigger_instances.')
        self.assertEqual(retrieved_trace_db.trigger_instances[0].object_id, trigger_instance_id,
                         'Expected updated trigger_instances.')

        Trace.delete(retrieved_trace_db)
        Trace.add_or_update(self.trace_empty)

def add_or_update_given_trace_db(trace_db, action_executions=None, rules=None,
                                 trigger_instances=None):
    """
    Will update an existing Trace.

    :param trace_db: The TraceDB to update.
    :type trace_db: ``TraceDB``

    :param action_executions: The action_execution to be added to the Trace. Should a list
                              of object_ids or a dict containing object_ids and caused_by.
    :type action_executions: ``list``

    :param rules: The rules to be added to the Trace. Should a list of object_ids or a dict
                  containing object_ids and caused_by.
    :type rules: ``list``

    :param trigger_instances: The trigger_instances to be added to the Trace. Should a list
                              of object_ids or a dict containing object_ids and caused_by.
    :type trigger_instances: ``list``

    :rtype: ``TraceDB``
    """
    if trace_db is None:
        raise ValueError('trace_db should be non-None.')

    if not action_executions:
        action_executions = []
    action_executions = [_to_trace_component_db(component=action_execution)
                         for action_execution in action_executions]

    if not rules:
        rules = []
    rules = [_to_trace_component_db(component=rule) for rule in rules]

    if not trigger_instances:
        trigger_instances = []
    trigger_instances = [_to_trace_component_db(component=trigger_instance)
                         for trigger_instance in trigger_instances]

    # If an id exists then this is an update and we do not want to perform
    # an upsert so use push_components which will use the push operator.
    if trace_db.id:
        return Trace.push_components(trace_db,
                                     action_executions=action_executions,
                                     rules=rules,
                                     trigger_instances=trigger_instances)

    trace_db.action_executions = action_executions
    trace_db.rules = rules
    trace_db.trigger_instances = trigger_instances

    return Trace.add_or_update(trace_db)

    def test_trace_provided(self):
        self.traceable_liveaction['context']['trace_context'] = {'trace_tag': 'OohLaLaLa'}
        action_services.request(self.traceable_liveaction)
        traces = Trace.get_all()
        self.assertEqual(len(traces), 1)
        self.assertEqual(len(traces[0]['action_executions']), 1)

        # Let's use existing trace id in trace context.
        # We shouldn't create new trace object.
        trace_id = str(traces[0].id)
        self.traceable_liveaction['context']['trace_context'] = {'id_': trace_id}
        action_services.request(self.traceable_liveaction)
        traces = Trace.get_all()
        self.assertEqual(len(traces), 1)
        self.assertEqual(len(traces[0]['action_executions']), 2)

 def test_trace_tag_resuse(self):
     self.traceable_liveaction['context']['trace_context'] = {'trace_tag': 'blank space'}
     action_services.request(self.traceable_liveaction)
     # Let's use same trace tag again and we should see two trace objects in db.
     action_services.request(self.traceable_liveaction)
     traces = Trace.query(**{'trace_tag': 'blank space'})
     self.assertEqual(len(traces), 2)

 def test_get(self):
     saved = TraceDBTest._create_save_trace(
         trace_tag='test_trace',
         action_executions=[str(bson.ObjectId()) for _ in range(4)],
         rules=[str(bson.ObjectId()) for _ in range(4)],
         trigger_instances=[str(bson.ObjectId()) for _ in range(5)])
     retrieved = Trace.get(id=saved.id)
     self.assertEquals(retrieved.id, saved.id, 'Incorrect trace retrieved.')

    def test_add_or_update_given_trace_context_new_with_causals(self):
        trace_context = {'trace_tag': 'causal_test_trace'}
        action_execution_id = 'action_execution_1'
        rule_id = 'rule_1'
        trigger_instance_id = 'trigger_instance_1'

        pre_add_or_update_traces = len(Trace.get_all())
        trace_db = trace_service.add_or_update_given_trace_context(
            trace_context,
            action_executions=[{'id': action_execution_id,
                                'caused_by': {'id': '%s:%s' % (rule_id, trigger_instance_id),
                                              'type': 'rule'}}],
            rules=[{'id': rule_id,
                    'caused_by': {'id': trigger_instance_id, 'type': 'trigger-instance'}}],
            trigger_instances=[trigger_instance_id])
        post_add_or_update_traces = len(Trace.get_all())

        self.assertTrue(post_add_or_update_traces > pre_add_or_update_traces,
                        'Expected new Trace to be created.')

        retrieved_trace_db = Trace.get_by_id(trace_db.id)
        self.assertEqual(len(retrieved_trace_db.action_executions), 1,
                         'Expected updated action_executions.')
        self.assertEqual(retrieved_trace_db.action_executions[0].object_id, action_execution_id,
                         'Expected updated action_executions.')
        self.assertEqual(retrieved_trace_db.action_executions[0].caused_by,
                         {'id': '%s:%s' % (rule_id, trigger_instance_id),
                          'type': 'rule'},
                         'Expected updated action_executions.')

        self.assertEqual(len(retrieved_trace_db.rules), 1, 'Expected updated rules.')
        self.assertEqual(retrieved_trace_db.rules[0].object_id, rule_id, 'Expected updated rules.')
        self.assertEqual(retrieved_trace_db.rules[0].caused_by,
                         {'id': trigger_instance_id, 'type': 'trigger-instance'},
                         'Expected updated rules.')

        self.assertEqual(len(retrieved_trace_db.trigger_instances), 1,
                         'Expected updated trigger_instances.')
        self.assertEqual(retrieved_trace_db.trigger_instances[0].object_id, trigger_instance_id,
                         'Expected updated trigger_instances.')
        self.assertEqual(retrieved_trace_db.trigger_instances[0].caused_by, {},
                         'Expected updated rules.')

        Trace.delete(retrieved_trace_db)

    def test_query(self):
        saved = TraceDBTest._create_save_trace(
            trace_tag='test_trace',
            action_executions=[str(bson.ObjectId()) for _ in range(4)],
            rules=[str(bson.ObjectId()) for _ in range(4)],
            trigger_instances=[str(bson.ObjectId()) for _ in range(5)])
        retrieved = Trace.query(trace_tag=saved.trace_tag)
        self.assertEquals(len(retrieved), 1, 'Should have 1 trace.')
        self.assertEquals(retrieved[0].id, saved.id, 'Incorrect trace retrieved.')

        # Add another trace with same trace_tag and confirm that we support.
        # This is most likley an anti-pattern for the trace_tag but it is an unknown.
        saved = TraceDBTest._create_save_trace(
            trace_tag='test_trace',
            action_executions=[str(bson.ObjectId()) for _ in range(2)],
            rules=[str(bson.ObjectId()) for _ in range(4)],
            trigger_instances=[str(bson.ObjectId()) for _ in range(3)])
        retrieved = Trace.query(trace_tag=saved.trace_tag)
        self.assertEquals(len(retrieved), 2, 'Should have 2 traces.')

def _get_single_trace_by_component(**component_filter):
    """
    Tries to return a single Trace mathing component_filter. Raises an exception
    when a filter matches multiple.
    """
    traces = Trace.query(**component_filter)
    if len(traces) == 0:
        return None
    elif len(traces) > 1:
        raise UniqueTraceNotFoundException(
            'More than 1 trace matching %s found.' % component_filter)
    return traces[0]

    def test_add_or_update_given_trace_db(self):
        action_execution_id = 'action_execution_1'
        rule_id = 'rule_1'
        trigger_instance_id = 'trigger_instance_1'
        to_save = copy.copy(self.trace_empty)
        to_save.id = None
        saved = trace_service.add_or_update_given_trace_db(
            to_save,
            action_executions=[action_execution_id],
            rules=[rule_id],
            trigger_instances=[trigger_instance_id])

        retrieved_trace_db = Trace.get_by_id(saved.id)
        self.assertEqual(len(retrieved_trace_db.action_executions), 1,
                         'Expected updated action_executions.')
        self.assertEqual(retrieved_trace_db.action_executions[0].object_id, action_execution_id,
                         'Expected updated action_executions.')

        self.assertEqual(len(retrieved_trace_db.rules), 1, 'Expected updated rules.')
        self.assertEqual(retrieved_trace_db.rules[0].object_id, rule_id, 'Expected updated rules.')

        self.assertEqual(len(retrieved_trace_db.trigger_instances), 1,
                         'Expected updated trigger_instances.')
        self.assertEqual(retrieved_trace_db.trigger_instances[0].object_id, trigger_instance_id,
                         'Expected updated trigger_instances.')

        # Now add more TraceComponents and validated that they are added properly.
        saved = trace_service.add_or_update_given_trace_db(
            retrieved_trace_db,
            action_executions=[str(bson.ObjectId()), str(bson.ObjectId())],
            rules=[str(bson.ObjectId())],
            trigger_instances=[str(bson.ObjectId()), str(bson.ObjectId()), str(bson.ObjectId())])
        retrieved_trace_db = Trace.get_by_id(saved.id)
        self.assertEqual(len(retrieved_trace_db.action_executions), 3,
                         'Expected updated action_executions.')
        self.assertEqual(len(retrieved_trace_db.rules), 2, 'Expected updated rules.')
        self.assertEqual(len(retrieved_trace_db.trigger_instances), 4,
                         'Expected updated trigger_instances.')

        Trace.delete(retrieved_trace_db)

def add_or_update_given_trace_db(trace_db, action_executions=None, rules=None,
                                 trigger_instances=None):
    """
    Will update an existing Trace.

    :param trace_db: The TraceDB to update.
    :type trace_db: ``TraceDB``

    :param action_executions: The action_execution to be added to the Trace. Should a list
                              of object_ids.
    :type action_executions: ``list``

    :param rules: The rules to be added to the Trace. Should a list of object_ids.
    :type rules: ``list``

    :param trigger_instances: The trigger_instances to be added to the Trace. Should a list
                              of object_ids.
    :type trigger_instances: ``list``

    :rtype: ``TraceDB``
    """
    if trace_db is None:
        raise ValueError('trace_db should be non-None.')

    if not action_executions:
        action_executions = []
    action_executions = [TraceComponentDB(object_id=action_execution)
                         for action_execution in action_executions]
    if trace_db.action_executions:
        action_executions.extend(trace_db.action_executions)

    if not rules:
        rules = []
    rules = [TraceComponentDB(object_id=rule) for rule in rules]
    if trace_db.rules:
        rules.extend(trace_db.rules)

    if not trigger_instances:
        trigger_instances = []
    trigger_instances = [TraceComponentDB(object_id=trigger_instance)
                         for trigger_instance in trigger_instances]
    if trace_db.trigger_instances:
        trigger_instances.extend(trace_db.trigger_instances)

    trace_db.action_executions = action_executions
    trace_db.rules = rules
    trace_db.trigger_instances = trigger_instances

    return Trace.add_or_update(trace_db)

def get_trace(trace_context, ignore_trace_tag=False):
    """
    :param trace_context: context object using which a trace can be found.
    :type trace_context: ``dict`` or ``TraceContext``

    :param ignore_trace_tag: Even if a trace_tag is provided will be ignored.
    :type ignore_trace_tag: ``str``

    :rtype: ``TraceDB``
    """

    trace_context = _get_valid_trace_context(trace_context)

    if not trace_context.id_ and not trace_context.trace_tag:
        raise ValueError('Atleast one of id_ or trace_tag should be specified.')

    if trace_context.id_:
        try:
            return Trace.get_by_id(trace_context.id_)
        except (ValidationError, ValueError):
            LOG.warning('Database lookup for Trace with id="%s" failed.',
                        trace_context.id_, exc_info=True)
            raise StackStormDBObjectNotFoundError(
                'Unable to find Trace with id="%s"' % trace_context.id_)

    if ignore_trace_tag:
        return None

    traces = Trace.query(trace_tag=trace_context.trace_tag)

    # Assume this method only handles 1 trace.
    if len(traces) > 1:
        raise UniqueTraceNotFoundException(
            'More than 1 Trace matching %s found.' % trace_context.trace_tag)

    return traces[0]

 def test_post_delete_trace(self):
     LIVE_ACTION_TRACE = copy.copy(LIVE_ACTION_1)
     LIVE_ACTION_TRACE['context'] = {'trace_context': {'trace_tag': 'balleilaka'}}
     post_resp = self._do_post(LIVE_ACTION_TRACE)
     self.assertEqual(post_resp.status_int, 201)
     delete_resp = self._do_delete(self._get_actionexecution_id(post_resp))
     self.assertEqual(delete_resp.status_int, 200)
     self.assertEqual(delete_resp.json['status'], 'canceled')
     trace_id = str(Trace.get_all()[0].id)
     LIVE_ACTION_TRACE['context'] = {'trace_context': {'id_': trace_id}}
     post_resp = self._do_post(LIVE_ACTION_TRACE)
     self.assertEqual(post_resp.status_int, 201)
     delete_resp = self._do_delete(self._get_actionexecution_id(post_resp))
     self.assertEqual(delete_resp.status_int, 200)
     self.assertEqual(delete_resp.json['status'], 'canceled')

 def test_post_delete_trace(self):
     """
     Validate that the API controller doesn't blow up on specifying
     trace_context.
     """
     LIVE_ACTION_TRACE = copy.copy(LIVE_ACTION_1)
     LIVE_ACTION_TRACE['context'] = {'trace_context': {'trace_tag': 'balleilaka'}}
     post_resp = self._do_post(LIVE_ACTION_TRACE)
     self.assertEqual(post_resp.status_int, 201)
     delete_resp = self._do_delete(self._get_actionexecution_id(post_resp))
     self.assertEqual(delete_resp.status_int, 200)
     self.assertEqual(delete_resp.json['status'], 'canceled')
     trace_id = str(Trace.get_all()[0].id)
     LIVE_ACTION_TRACE['context'] = {'trace_context': {'id_': trace_id}}
     post_resp = self._do_post(LIVE_ACTION_TRACE)
     self.assertEqual(post_resp.status_int, 201)
     delete_resp = self._do_delete(self._get_actionexecution_id(post_resp))
     self.assertEqual(delete_resp.status_int, 200)
     self.assertEqual(delete_resp.json['status'], 'canceled')

    def test_update_via_list_push_components(self):
        no_action_executions = 4
        no_rules = 4
        no_trigger_instances = 5
        saved = TraceDBTest._create_save_trace(
            trace_tag='test_trace',
            action_executions=[str(bson.ObjectId()) for _ in range(no_action_executions)],
            rules=[str(bson.ObjectId()) for _ in range(no_rules)],
            trigger_instances=[str(bson.ObjectId()) for _ in range(no_trigger_instances)])

        retrieved = Trace.push_components(
            saved,
            action_executions=[TraceComponentDB(object_id=str(bson.ObjectId()))
                               for _ in range(no_action_executions)],
            rules=[TraceComponentDB(object_id=str(bson.ObjectId()))
                   for _ in range(no_rules)],
            trigger_instances=[TraceComponentDB(object_id=str(bson.ObjectId()))
                               for _ in range(no_trigger_instances)])

        self.assertEquals(retrieved.id, saved.id, 'Incorrect trace retrieved.')
        self.assertEquals(len(retrieved.action_executions), no_action_executions * 2)
        self.assertEquals(len(retrieved.rules), no_rules * 2)
        self.assertEquals(len(retrieved.trigger_instances), no_trigger_instances * 2)

    def _create_save_trace(trace_tag, id_=None, action_executions=None, rules=None,
                           trigger_instances=None):

        if action_executions is None:
            action_executions = []
        action_executions = [TraceComponentDB(object_id=action_execution)
                             for action_execution in action_executions]

        if rules is None:
            rules = []
        rules = [TraceComponentDB(object_id=rule) for rule in rules]

        if trigger_instances is None:
            trigger_instances = []
        trigger_instances = [TraceComponentDB(object_id=trigger_instance)
                             for trigger_instance in trigger_instances]

        created = TraceDB(id=id_,
                          trace_tag=trace_tag,
                          trigger_instances=trigger_instances,
                          rules=rules,
                          action_executions=action_executions)
        return Trace.add_or_update(created)

def get_trace_db_by_rule(rule=None, rule_id=None):
    if rule:
        rule_id = str(rule.id)
    # by rule could return multiple traces
    return Trace.query(rules__object_id=rule_id)