diff options
Diffstat (limited to 'distbuild/mainloop.py')
-rw-r--r-- | distbuild/mainloop.py | 65 |
1 files changed, 65 insertions, 0 deletions
diff --git a/distbuild/mainloop.py b/distbuild/mainloop.py index e7c0cc3b..d48d60ec 100644 --- a/distbuild/mainloop.py +++ b/distbuild/mainloop.py @@ -132,3 +132,68 @@ class MainLoop(object): event_source, event = self._events.pop(0) yield event_source, event + + +class TestableMainLoop(MainLoop): + '''Special mainloop class with extra hooks for tests to use. + + When writing a test, you often need to wait until a certain event has + happened, then examine that event. The run_until_event() and + run_until_new_state_machine() functions allow this. + + ''' + def __init__(self): + super(TestableMainLoop, self).__init__() + + self._machines_added_this_cycle = [] + self._events_sent_this_cycle = [] + + def add_state_machine(self, machine): + # Overriding the base class to monitor new state machines. + super(TestableMainLoop, self).add_state_machine(machine) + self._machines_added_this_cycle.append(machine) + + def queue_event(self, event_source, event): + # Overriding the base class to monitor new events. + super(TestableMainLoop, self).queue_event(event_source, event) + self._events_sent_this_cycle.append((event_source, event)) + + def run_until_event(self, target_event_source, target_event_type): + '''Run the main loop continuously until a given event happens. + + All queued messages will be processed before the loop exits. + + ''' + logging.debug('Running main loop until a %s event from %s.', + target_event_type, target_event_source) + while self._machines: + self._events_sent_this_cycle = [] + + self._run_once() + + for event_source, event in self._events_sent_this_cycle: + if target_event_source == event_source: + if isinstance(event, target_event_type): + logging.debug( + 'Received %s from %s, exiting loop.', event, + event_source) + return event + + def run_until_new_state_machine(self, target_machine_type): + '''Run the main loop continuously until a new state machine appears. + + All queued messages will be processed before the loop exits. + + ''' + logging.debug('Running main loop until a new %s appears.', + target_machine_type) + while self._machines: + self._machines_added_this_cycle = [] + + self._run_once() + + for machine in self._machines_added_this_cycle: + if type(machine) == target_machine_type: + logging.debug( + 'Found new machine %s, exiting loop.', machine) + return machine |