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)