def _SetupUserNotificationTimerangeTest(self):
d = self.db
username = "******"
d.WriteGRRUser(username)
ts = []
ts.append(rdfvalue.RDFDatetime.Now())
n = objects.UserNotification(
username=username,
notification_type=objects.UserNotification.Type.
TYPE_CLIENT_INTERROGATED,
state=objects.UserNotification.State.STATE_PENDING,
message="n0")
d.WriteUserNotification(n)
ts.append(rdfvalue.RDFDatetime.Now())
n = objects.UserNotification(
username=username,
notification_type=objects.UserNotification.Type.
TYPE_CLIENT_INTERROGATED,
state=objects.UserNotification.State.STATE_PENDING,
message="n1")
d.WriteUserNotification(n)
ts.append(rdfvalue.RDFDatetime.Now())
return ts
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 = objects.UserNotification(
username=username,
notification_type=objects.UserNotification.Type.
TYPE_CLIENT_INTERROGATED,
state=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 testNotificationTimestampIsGeneratedWhenNotExplicit(self):
d = self.db
username = "******"
d.WriteGRRUser(username)
n = objects.UserNotification(
username=username,
notification_type=objects.UserNotification.Type.
TYPE_CLIENT_INTERROGATED,
state=objects.UserNotification.State.STATE_PENDING,
message="blah")
d.WriteUserNotification(n)
ns = d.ReadUserNotifications(username)
self.assertEqual(len(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 = [
objects.UserNotification(
username=username,
notification_type=objects.UserNotification.Type.
TYPE_CLIENT_INTERROGATED,
state=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))