def TestHuntHelperWithMultipleMocks(client_mocks,
check_flow_errors=False,
token=None,
iteration_limit=None):
"""Runs a hunt with a given set of clients mocks.
Args:
client_mocks: Dictionary of (client_id->client_mock) pairs. Client mock
objects are used to handle client actions. Methods names of a client mock
object correspond to client actions names. For an example of a client mock
object, see SampleHuntMock.
check_flow_errors: If True, raises when one of hunt-initiated flows fails.
token: An instance of access_control.ACLToken security token.
iteration_limit: If None, hunt will run until it's finished. Otherwise,
worker_mock.Next() will be called iteration_limit number of times. Every
iteration processes worker's message queue. If new messages are sent to
the queue during the iteration processing, they will be processed on next
iteration,
"""
total_flows = set()
# Worker always runs with absolute privileges, therefore making the token
# SetUID().
token = token.SetUID()
client_mocks = [
flow_test_lib.MockClient(client_id, client_mock, token=token)
for client_id, client_mock in iteritems(client_mocks)
]
with flow_test_lib.TestWorker(token=True) as rel_db_worker:
worker_mock = worker_test_lib.MockWorker(
check_flow_errors=check_flow_errors, token=token)
# Run the clients and worker until nothing changes any more.
while iteration_limit is None or iteration_limit > 0:
client_processed = 0
for client_mock in client_mocks:
client_processed += client_mock.Next()
flows_run = []
for flow_run in worker_mock.Next():
total_flows.add(flow_run)
flows_run.append(flow_run)
worker_processed = rel_db_worker.ResetProcessedFlows()
flows_run.extend(worker_processed)
if client_processed == 0 and not flows_run:
break
if iteration_limit:
iteration_limit -= 1
if check_flow_errors:
flow_test_lib.CheckFlowErrors(total_flows, token=token)
def TestHuntHelperWithMultipleMocks(client_mocks,
check_flow_errors=False,
token=None,
iteration_limit=None,
worker=None):
"""Runs a hunt with a given set of clients mocks.
Args:
client_mocks: Dictionary of (client_id->client_mock) pairs. Client mock
objects are used to handle client actions. Methods names of a client mock
object correspond to client actions names. For an example of a client mock
object, see SampleHuntMock.
check_flow_errors: If True, raises when one of hunt-initiated flows fails.
token: An instance of access_control.ACLToken security token.
iteration_limit: If None, hunt will run until it's finished. Otherwise,
worker_mock.Next() will be called iteration_limit number of times. Every
iteration processes worker's message queue. If new messages are sent to
the queue during the iteration processing, they will be processed on next
iteration.
worker: flow_test_lib.TestWorker object to use.
Returns:
A number of iterations complete.
"""
if token is None:
token = access_control.ACLToken(username="******")
total_flows = set()
# Worker always runs with absolute privileges, therefore making the token
# SetUID().
token = token.SetUID()
client_mocks = [
flow_test_lib.MockClient(client_id, client_mock, token=token)
for client_id, client_mock in iteritems(client_mocks)
]
if worker is None:
rel_db_worker = flow_test_lib.TestWorker(threadpool_size=0, token=True)
data_store.REL_DB.RegisterFlowProcessingHandler(rel_db_worker.ProcessFlow)
else:
rel_db_worker = worker
num_iterations = 0
try:
worker_mock = worker_test_lib.MockWorker(
check_flow_errors=check_flow_errors, token=token)
# Run the clients and worker until nothing changes any more.
while iteration_limit is None or num_iterations < iteration_limit:
worker_processed = []
if data_store.RelationalDBEnabled():
data_store.REL_DB.delegate.WaitUntilNoFlowsToProcess(timeout=10)
worker_processed = rel_db_worker.ResetProcessedFlows()
client_processed = 0
for client_mock in client_mocks:
client_processed += client_mock.Next()
flows_run = []
for flow_run in worker_mock.Next():
total_flows.add(flow_run)
flows_run.append(flow_run)
flows_run.extend(worker_processed)
num_iterations += 1
if client_processed == 0 and not flows_run and not worker_processed:
break
if check_flow_errors:
flow_test_lib.CheckFlowErrors(total_flows, token=token)
finally:
if worker is None:
data_store.REL_DB.UnregisterFlowProcessingHandler(timeout=60)
rel_db_worker.Shutdown()
return num_iterations