def testSizeLimitedQueue(self):
queue = comms.SizeLimitedQueue(maxsize=10000000, heart_beat_cb=lambda: None)
msg_a = rdf_flows.GrrMessage(name="A")
msg_b = rdf_flows.GrrMessage(name="B")
msg_c = rdf_flows.GrrMessage(name="C")
for _ in xrange(10):
queue.Put(msg_a, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY)
queue.Put(msg_b, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY)
queue.Put(msg_c, rdf_flows.GrrMessage.Priority.HIGH_PRIORITY)
result = queue.GetMessages()
self.assertEqual(list(result.job), [msg_c] * 10 + [msg_a, msg_b] * 10)
# Tests a partial Get().
for _ in xrange(7):
queue.Put(msg_a, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY)
queue.Put(msg_b, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY)
queue.Put(msg_c, rdf_flows.GrrMessage.Priority.HIGH_PRIORITY)
result = queue.GetMessages(
soft_size_limit=len(msg_a.SerializeToString()) * 5 - 1)
self.assertEqual(list(result.job), [msg_c] * 5)
for _ in xrange(3):
queue.Put(msg_a, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY)
queue.Put(msg_b, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY)
queue.Put(msg_c, rdf_flows.GrrMessage.Priority.HIGH_PRIORITY)
# Append the remaining messages to the same result.
result.job.Extend(queue.GetMessages().job)
self.assertEqual(list(result.job), [msg_c] * 10 + [msg_a, msg_b] * 10)
def testSizeLimitedQueueSize(self):
q = comms.SizeLimitedQueue(1000)
msg_a = rdf_flows.GrrMessage(name="A")
msg_b = rdf_flows.GrrMessage(name="B")
msg_c = rdf_flows.GrrMessage(name="C")
msg_d = rdf_flows.GrrMessage(name="D")
messages = [msg_a, msg_b, msg_c, msg_d]
in_queue = []
self.assertEqual(q.Size(), 0)
for m in messages:
q.Put(m, block=False)
in_queue.append(m)
self.assertEqual(
q.Size(), sum([len(m.SerializeToBytes()) for m in in_queue]))
for _ in range(len(messages)):
msg_list = q.GetMessages(1)
self.assertLen(msg_list.job, 1)
in_queue.remove(msg_list.job[0])
self.assertEqual(
q.Size(), sum([len(m.SerializeToBytes()) for m in in_queue]))
def testSizeLimitedQueueOverflow(self):
msg_a = rdf_flows.GrrMessage(name="A")
msg_b = rdf_flows.GrrMessage(name="B")
msg_c = rdf_flows.GrrMessage(name="C")
msg_d = rdf_flows.GrrMessage(name="D")
limited_queue = comms.SizeLimitedQueue(
maxsize=3 * len(msg_a.SerializeToString()), heart_beat_cb=lambda: None)
limited_queue.Put(msg_a, block=False)
limited_queue.Put(msg_b, block=False)
limited_queue.Put(msg_c, block=False)
with self.assertRaises(queue.Full):
limited_queue.Put(msg_d, block=False)
def testSizeLimitedQueueOverflow(self):
msg_a = rdf_flows.GrrMessage(name="A")
msg_b = rdf_flows.GrrMessage(name="B")
msg_c = rdf_flows.GrrMessage(name="C")
msg_d = rdf_flows.GrrMessage(name="D")
queue = comms.SizeLimitedQueue(
maxsize=3 * len(msg_a.SerializeToString()), heart_beat_cb=lambda: None)
queue.Put(msg_a, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, block=False)
queue.Put(msg_b, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, block=False)
queue.Put(msg_c, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, block=False)
with self.assertRaises(Queue.Full):
queue.Put(
msg_d, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, block=False)
def testSizeLimitedQueueHeartbeat(self):
msg_a = rdf_flows.GrrMessage(name="A")
msg_b = rdf_flows.GrrMessage(name="B")
msg_c = rdf_flows.GrrMessage(name="C")
msg_d = rdf_flows.GrrMessage(name="D")
heartbeat = mock.Mock()
queue = comms.SizeLimitedQueue(
maxsize=3 * len(msg_a.SerializeToString()), heart_beat_cb=heartbeat)
queue.Put(msg_a, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY)
queue.Put(msg_b, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY)
queue.Put(msg_c, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY)
with self.assertRaises(Queue.Full):
queue.Put(msg_d, rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, timeout=1)
self.assertTrue(heartbeat.called)
def testSizeLimitedQueueHeartbeat(self):
msg_a = rdf_flows.GrrMessage(name="A")
msg_b = rdf_flows.GrrMessage(name="B")
msg_c = rdf_flows.GrrMessage(name="C")
msg_d = rdf_flows.GrrMessage(name="D")
heartbeat = mock.Mock()
limited_queue = comms.SizeLimitedQueue(
maxsize=3 * len(msg_a.SerializeToString()), heart_beat_cb=heartbeat)
limited_queue.Put(msg_a)
limited_queue.Put(msg_b)
limited_queue.Put(msg_c)
with self.assertRaises(queue.Full):
limited_queue.Put(msg_d, timeout=1)
self.assertTrue(heartbeat.called)
def testSizeLimitedQueue(self):
limited_queue = comms.SizeLimitedQueue(maxsize=10000000,
heart_beat_cb=lambda: None)
msg_a = rdf_flows.GrrMessage(name="A")
msg_b = rdf_flows.GrrMessage(name="B")
msg_c = rdf_flows.GrrMessage(name="C")
for _ in range(10):
limited_queue.Put(msg_a)
limited_queue.Put(msg_b)
limited_queue.Put(msg_c)
result = limited_queue.GetMessages()
self.assertCountEqual(list(result.job),
[msg_c] * 10 + [msg_a, msg_b] * 10)
# Tests a partial Get().
for _ in range(7):
limited_queue.Put(msg_a)
limited_queue.Put(msg_b)
limited_queue.Put(msg_c)
result = limited_queue.GetMessages(
soft_size_limit=len(msg_a.SerializeToString()) * 5 - 1)
self.assertLen(list(result.job), 5)
for _ in range(3):
limited_queue.Put(msg_a)
limited_queue.Put(msg_b)
limited_queue.Put(msg_c)
# Append the remaining messages to the same result.
result.job.Extend(limited_queue.GetMessages().job)
self.assertCountEqual(list(result.job),
[msg_c] * 10 + [msg_a, msg_b] * 10)