def InitTestServer(self):
prefix = "pool-%s" % self._testMethodName
self.server = front_end.FrontEndServer(
certificate=config.CONFIG["Frontend.certificate"],
private_key=config.CONFIG["PrivateKeys.server_key"],
message_expiry_time=self.message_expiry_time,
threadpool_prefix=prefix)
def __init__(self,
server_address,
handler,
frontend=None,
*args,
**kwargs):
stats.STATS.SetGaugeValue("frontend_max_active_count",
self.request_queue_size)
if frontend:
self.frontend = frontend
else:
self.frontend = front_end.FrontEndServer(
certificate=config.CONFIG["Frontend.certificate"],
private_key=config.CONFIG["PrivateKeys.server_key"],
max_queue_size=config.CONFIG["Frontend.max_queue_size"],
message_expiry_time=config.
CONFIG["Frontend.message_expiry_time"],
max_retransmission_time=config.
CONFIG["Frontend.max_retransmission_time"])
self.server_cert = config.CONFIG["Frontend.certificate"]
(address, _) = server_address
version = ipaddr.IPAddress(address).version
if version == 4:
self.address_family = socket.AF_INET
elif version == 6:
self.address_family = socket.AF_INET6
logging.info("Will attempt to listen on %s", server_address)
BaseHTTPServer.HTTPServer.__init__(self, server_address, handler,
*args, **kwargs)
def __init__(self):
self.frontend = front_end.FrontEndServer(
certificate=config.CONFIG["Frontend.certificate"],
private_key=config.CONFIG["PrivateKeys.server_key"],
max_queue_size=config.CONFIG["Frontend.max_queue_size"],
message_expiry_time=config.CONFIG["Frontend.message_expiry_time"],
max_retransmission_time=config.CONFIG[
"Frontend.max_retransmission_time"])
def testEqualTimestampNotifications(self):
frontend_server = front_end.FrontEndServer(
certificate=config.CONFIG["Frontend.certificate"],
private_key=config.CONFIG["PrivateKeys.server_key"],
message_expiry_time=100,
threadpool_prefix="notification-test")
# This schedules 10 requests.
session_id = flow.GRRFlow.StartFlow(
client_id=self.client_id,
flow_name="WorkerSendingTestFlow",
token=self.token)
# We pretend that the client processed all the 10 requests at once and
# sends the replies in a single http poll.
messages = [
rdf_flows.GrrMessage(
request_id=i,
response_id=1,
session_id=session_id,
payload=rdf_protodict.DataBlob(string="test%s" % i),
generate_task_id=True) for i in range(1, 11)
]
status = rdf_flows.GrrStatus(status=rdf_flows.GrrStatus.ReturnedStatus.OK)
statuses = [
rdf_flows.GrrMessage(
request_id=i,
response_id=2,
session_id=session_id,
payload=status,
type=rdf_flows.GrrMessage.Type.STATUS,
generate_task_id=True) for i in range(1, 11)
]
frontend_server.ReceiveMessages(self.client_id, messages + statuses)
with queue_manager.QueueManager(token=self.token) as q:
all_notifications = q.GetNotificationsByPriorityForAllShards(
rdfvalue.RDFURN("aff4:/F"))
medium_priority = rdf_flows.GrrNotification.Priority.MEDIUM_PRIORITY
medium_notifications = all_notifications[medium_priority]
my_notifications = [
n for n in medium_notifications if n.session_id == session_id
]
# There must not be more than one notification.
self.assertEqual(len(my_notifications), 1)
notification = my_notifications[0]
self.assertEqual(notification.first_queued, notification.timestamp)
self.assertEqual(notification.last_status, 10)
