def testUpdateUserNotificationsNoTimestamps(self): self.db.WriteGRRUser("foo") notification = rdf_objects.UserNotification( username="******", message="Lorem ipsum.", state=rdf_objects.UserNotification.State.STATE_PENDING, notification_type=( rdf_objects.UserNotification.Type.TYPE_CLIENT_INTERROGATED)) self.db.WriteUserNotification(notification) notification = rdf_objects.UserNotification( username="******", message="Dolor sit amet.", state=rdf_objects.UserNotification.State.STATE_PENDING, notification_type=( rdf_objects.UserNotification.Type.TYPE_CLIENT_INTERROGATED)) self.db.WriteUserNotification(notification) # Should not raise and should not change anything. self.db.UpdateUserNotifications( username="******", timestamps=[], state=rdf_objects.UserNotification.State.STATE_NOT_PENDING) notifications = self.db.ReadUserNotifications(username="******") self.assertLen(notifications, 2) self.assertEqual(notifications[0].state, rdf_objects.UserNotification.State.STATE_PENDING) self.assertEqual(notifications[1].state, rdf_objects.UserNotification.State.STATE_PENDING)
def testMultipleNotificationsCanBeWrittenAndRead(self): username = "******" self.db.WriteGRRUser(username) # pylint: disable=invalid-name NotificationType = rdf_objects.UserNotification.Type NotificationState = rdf_objects.UserNotification.State # pylint: enable=invalid-name self.db.WriteUserNotification( rdf_objects.UserNotification( username=username, notification_type=NotificationType.TYPE_CLIENT_INTERROGATED, state=NotificationState.STATE_PENDING, message="Lorem ipsum.")) self.db.WriteUserNotification( rdf_objects.UserNotification( username=username, notification_type=NotificationType. TYPE_CLIENT_APPROVAL_REQUESTED, state=NotificationState.STATE_NOT_PENDING, message="Dolor sit amet.")) self.db.WriteUserNotification( rdf_objects.UserNotification( username=username, notification_type=NotificationType.TYPE_FLOW_RUN_FAILED, state=NotificationState.STATE_PENDING, message="Consectetur adipiscing elit.")) notifications = self.db.ReadUserNotifications(username) # TODO(hanuszczak): Database API should guaranee to return notifications # ordered by timestamp. notifications.sort(key=lambda notification: notification.timestamp) self.assertLen(notifications, 3) self.assertEqual(notifications[0].username, username) self.assertEqual(notifications[0].notification_type, NotificationType.TYPE_CLIENT_INTERROGATED) self.assertEqual(notifications[0].state, NotificationState.STATE_PENDING) self.assertEqual(notifications[0].message, "Lorem ipsum.") self.assertEqual(notifications[1].username, username) self.assertEqual(notifications[1].notification_type, NotificationType.TYPE_CLIENT_APPROVAL_REQUESTED) self.assertEqual(notifications[1].state, NotificationState.STATE_NOT_PENDING) self.assertEqual(notifications[1].message, "Dolor sit amet.") self.assertEqual(notifications[2].username, username) self.assertEqual(notifications[2].notification_type, NotificationType.TYPE_FLOW_RUN_FAILED) self.assertEqual(notifications[2].state, NotificationState.STATE_PENDING) self.assertEqual(notifications[2].message, "Consectetur adipiscing elit.")
def _SetupUserNotificationTimerangeTest(self, username="******"): d = self.db d.WriteGRRUser(username) ts = [] ts.append(rdfvalue.RDFDatetime.Now()) n = rdf_objects.UserNotification( username=username, notification_type=rdf_objects.UserNotification.Type. TYPE_CLIENT_INTERROGATED, state=rdf_objects.UserNotification.State.STATE_PENDING, message="n0") d.WriteUserNotification(n) ts.append(rdfvalue.RDFDatetime.Now()) n = rdf_objects.UserNotification( username=username, notification_type=rdf_objects.UserNotification.Type. TYPE_CLIENT_INTERROGATED, state=rdf_objects.UserNotification.State.STATE_PENDING, message="n1") d.WriteUserNotification(n) ts.append(rdfvalue.RDFDatetime.Now()) return ts
def testNotificationForUnknownUser(self): n = rdf_objects.UserNotification( username="******", notification_type=rdf_objects.UserNotification.Type. TYPE_CLIENT_INTERROGATED, state=rdf_objects.UserNotification.State.STATE_PENDING) with self.assertRaises(db.UnknownGRRUserError): self.db.WriteUserNotification(n)
def _Notify(username, notification_type, message, object_reference): """Schedules a new-style REL_DB user notification.""" if object_reference: uc = object_reference.UnionCast() if hasattr(uc, "client_id"): message = _HostPrefix(uc.client_id) + message n = rdf_objects.UserNotification( username=username, notification_type=notification_type, state=rdf_objects.UserNotification.State.STATE_PENDING, message=message, reference=object_reference) data_store.REL_DB.WriteUserNotification(n)
def testNotificationCanBeWrittenAndRead(self): d = self.db username = "******" d.WriteGRRUser(username) n = rdf_objects.UserNotification( username=username, notification_type=rdf_objects.UserNotification.Type. TYPE_CLIENT_INTERROGATED, state=rdf_objects.UserNotification.State.STATE_PENDING, timestamp=rdfvalue.RDFDatetime(42), message="blah") d.WriteUserNotification(n) ns = d.ReadUserNotifications(username) self.assertEqual(len(ns), 1) self.assertEqual(ns[0], n)
def testNotificationCanBeWrittenAndRead(self): d = self.db username = "******" d.WriteGRRUser(username) n = rdf_objects.UserNotification( username=username, notification_type=rdf_objects.UserNotification.Type. TYPE_CLIENT_INTERROGATED, state=rdf_objects.UserNotification.State.STATE_PENDING, message="blah") d.WriteUserNotification(n) ns = d.ReadUserNotifications(username) self.assertLen(ns, 1) ns[0].timestamp = None # Database generates timestamps, not interesting. self.assertEqual(ns[0], n)
def Notify(username, notification_type, message, object_reference): """Schedules a new-style REL_DB user notification.""" # Do not try to notify system users (e.g. Cron). if username in access_control.SYSTEM_USERS: return if object_reference: uc = object_reference.UnionCast() if hasattr(uc, "client_id"): message = _HostPrefix(uc.client_id) + message n = rdf_objects.UserNotification( username=username, notification_type=notification_type, state=rdf_objects.UserNotification.State.STATE_PENDING, message=message, reference=object_reference) data_store.REL_DB.WriteUserNotification(n)
def testNotificationTimestampIsGeneratedWhenNotExplicit(self): d = self.db username = "******" d.WriteGRRUser(username) n = rdf_objects.UserNotification( username=username, notification_type=rdf_objects.UserNotification.Type. TYPE_CLIENT_INTERROGATED, state=rdf_objects.UserNotification.State.STATE_PENDING, message="blah") d.WriteUserNotification(n) ns = d.ReadUserNotifications(username) self.assertLen(ns, 1) self.assertNotEqual(int(ns[0].timestamp), 0) self.assertNotEqual(ns[0], n) n.timestamp = ns[0].timestamp self.assertEqual(ns[0], n)
def testMultipleNotificationsCanBeWrittenAndRead(self): d = self.db username = "******" d.WriteGRRUser(username) ns = [ rdf_objects.UserNotification( username=username, notification_type=rdf_objects.UserNotification.Type. TYPE_CLIENT_INTERROGATED, state=rdf_objects.UserNotification.State.STATE_PENDING, timestamp=rdfvalue.RDFDatetime(42 + i), message="blah%d" % i) for i in range(10) ] for n in ns: d.WriteUserNotification(n) read_ns = d.ReadUserNotifications(username) self.assertEqual(len(read_ns), 10) self.assertEqual(ns, sorted(read_ns, key=lambda x: x.timestamp))