action.Execute(message)
self.action_counts[message.name] += 1
return self.client_worker.Drain()
@property
def recorded_messages(self):
return self._recorded_messages
class Store(server_stubs.ClientActionStub):
"""A test client action."""
pass
action_registry.RegisterAdditionalTestClientAction(Store)
class CPULimitClientMock(ActionMock):
"""A mock for testing resource limits."""
in_rdfvalue = rdf_protodict.DataBlob
out_rdfvalues = []
def __init__(self,
storage=None,
user_cpu_usage=None,
system_cpu_usage=None,
network_usage=None,
runtime_us=None):
super(CPULimitClientMock, self).__init__()
raise parsers.ParseError("It was bound to happen.")
# TODO: These should be defined next to the test they are used in
# once the metaclass registry madness is resolved.
# pylint: disable=function-redefined
class FooAction(server_stubs.ClientActionStub):
in_rdfvalue = None
out_rdfvalues = [rdfvalue.RDFString]
action_registry.RegisterAdditionalTestClientAction(FooAction)
class FooAction(actions.ActionPlugin):
in_rdfvalue = None
out_rdfvalues = [rdfvalue.RDFString]
def Run(self, args):
del args # Unused.
self.SendReply(rdfvalue.RDFString("zażółć gęślą jaźń 🎮"))
# pylint: enable=function-redefined
from grr.test_lib import hunt_test_lib
from grr.test_lib import notification_test_lib
from grr.test_lib import test_lib
from grr.test_lib import test_output_plugins
class ReturnHello(actions.ActionPlugin):
"""A test client action."""
out_rdfvalues = [rdfvalue.RDFString]
def Run(self, _):
self.SendReply(rdfvalue.RDFString("Hello World"))
action_registry.RegisterAdditionalTestClientAction(ReturnHello)
class ClientMock(action_mocks.ActionMock):
"""Mock of client actions."""
def __init__(self):
super(ClientMock, self).__init__(ReturnHello)
class CallStateFlow(flow_base.FlowBase):
"""A flow that calls one of its own states."""
# This is a global flag which will be set when the flow runs.
success = False
