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))
