def testPreparePredefinedMessageForNote_withNewContent(self):
    lab_name = "alab"
    content = "content1"

    message = note_manager.PreparePredefinedMessageForNote(
        api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON,
        lab_name=lab_name,
        content=content)

    self.assertIsNone(message.key)  # new entity without a key yet
    self.assertEqual(lab_name, message.lab_name)
    self.assertEqual(content, message.content)
    self.assertEqual(1, message.used_count)
  def testPreparePredefinedMessageForNote_withValidId(self):
    message_id = 111
    lab_name = "alab"
    content = "content1"
    datastore_entities.PredefinedMessage(
        key=ndb.Key(datastore_entities.PredefinedMessage, message_id),
        lab_name=lab_name,
        type=api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON,
        content=content,
        used_count=2).put()

    message = note_manager.PreparePredefinedMessageForNote(
        api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON,
        message_id=message_id)

    self.assertEqual(message_id, message.key.id())
    self.assertEqual(lab_name, message.lab_name)
    self.assertEqual(content, message.content)
    self.assertEqual(3, message.used_count)  # the used_count increases
Esempio n. 3
0
  def BatchUpdateNotesWithPredefinedMessage(self, request):
    """Batch update notes with the same predefined message.

    Args:
      request: an API request.

    Returns:
      an api_messages.NoteCollection object.
    """
    time_now = datetime.datetime.utcnow()

    host_note_entities = []
    for note in request.notes:
      note_id = int(note.id) if note.id is not None else None
      host_note_entity = datastore_util.GetOrCreateEntity(
          datastore_entities.Note,
          entity_id=note_id,
          hostname=note.hostname,
          type=common.NoteType.HOST_NOTE)
      host_note_entity.populate(
          user=request.user,
          message=request.message,
          timestamp=time_now,
          event_time=request.event_time)
      host_note_entities.append(host_note_entity)

    try:
      offline_reason_entity = note_manager.PreparePredefinedMessageForNote(
          common.PredefinedMessageType.HOST_OFFLINE_REASON,
          message_id=request.offline_reason_id,
          lab_name=request.lab_name,
          content=request.offline_reason,
          delta_count=len(host_note_entities))
    except note_manager.InvalidParameterError as err:
      raise endpoints.BadRequestException("Invalid offline reason: [%s]" % err)
    if offline_reason_entity:
      for host_note_entity in host_note_entities:
        host_note_entity.offline_reason = offline_reason_entity.content
      offline_reason_entity.put()

    try:
      recovery_action_entity = note_manager.PreparePredefinedMessageForNote(
          common.PredefinedMessageType.HOST_RECOVERY_ACTION,
          message_id=request.recovery_action_id,
          lab_name=request.lab_name,
          content=request.recovery_action,
          delta_count=len(host_note_entities))
    except note_manager.InvalidParameterError as err:
      raise endpoints.BadRequestException("Invalid recovery action: [%s]" % err)
    if recovery_action_entity:
      for host_note_entity in host_note_entities:
        host_note_entity.recovery_action = recovery_action_entity.content
      recovery_action_entity.put()

    note_keys = ndb.put_multi(host_note_entities)
    host_note_entities = ndb.get_multi(note_keys)
    note_msgs = []
    for host_note_entity in host_note_entities:
      host_note_msg = datastore_entities.ToMessage(host_note_entity)
      note_msgs.append(host_note_msg)

      host_note_event_msg = api_messages.NoteEvent(
          note=host_note_msg,
          lab_name=request.lab_name)
      note_manager.PublishMessage(
          host_note_event_msg, common.PublishEventType.HOST_NOTE_EVENT)

    for request_note, updated_note_key in zip(request.notes, note_keys):
      if not request_note.id:
        # If ids are not provided, then a new note is created, we should create
        # a history snapshot.
        device_manager.CreateAndSaveHostInfoHistoryFromHostNote(
            request_note.hostname, updated_note_key.id())

    return api_messages.NoteCollection(
        notes=note_msgs, more=False, next_cursor=None, prev_cursor=None)
 def testPreparePredefinedMessageForNote_withInvalidId(self):
   invalid_message_id = 111
   with self.assertRaises(note_manager.InvalidParameterError):
     note_manager.PreparePredefinedMessageForNote(
         api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON,
         message_id=invalid_message_id)
Esempio n. 5
0
  def AddOrUpdateNote(self, request):
    """Add or update a host note.

    Args:
      request: an API request.

    Returns:
      an api_messages.Note.
    """
    time_now = datetime.datetime.utcnow()

    host_note_entity = datastore_util.GetOrCreateEntity(
        datastore_entities.Note,
        entity_id=request.id,
        hostname=request.hostname,
        type=common.NoteType.HOST_NOTE)
    host_note_entity.populate(
        user=request.user,
        message=request.message,
        timestamp=time_now,
        event_time=request.event_time)
    entities_to_update = [host_note_entity]

    try:
      offline_reason_entity = note_manager.PreparePredefinedMessageForNote(
          common.PredefinedMessageType.HOST_OFFLINE_REASON,
          message_id=request.offline_reason_id,
          lab_name=request.lab_name,
          content=request.offline_reason)
    except note_manager.InvalidParameterError as err:
      raise endpoints.BadRequestException("Invalid offline reason: [%s]" % err)
    if offline_reason_entity:
      host_note_entity.offline_reason = offline_reason_entity.content
      entities_to_update.append(offline_reason_entity)

    try:
      recovery_action_entity = note_manager.PreparePredefinedMessageForNote(
          common.PredefinedMessageType.HOST_RECOVERY_ACTION,
          message_id=request.recovery_action_id,
          lab_name=request.lab_name,
          content=request.recovery_action)
    except note_manager.InvalidParameterError as err:
      raise endpoints.BadRequestException("Invalid recovery action: [%s]" % err)
    if recovery_action_entity:
      host_note_entity.recovery_action = recovery_action_entity.content
      entities_to_update.append(recovery_action_entity)

    keys = ndb.put_multi(entities_to_update)
    host_note_msg = datastore_entities.ToMessage(host_note_entity)

    host_note_event_msg = api_messages.NoteEvent(
        note=host_note_msg, lab_name=request.lab_name)
    note_manager.PublishMessage(host_note_event_msg,
                                common.PublishEventType.HOST_NOTE_EVENT)

    note_key = keys[0]
    if request.id != note_key.id():
      # If ids are different, then a new note is created, we should create
      # a history snapshot.
      device_manager.CreateAndSaveHostInfoHistoryFromHostNote(
          request.hostname, note_key.id())

    return host_note_msg