1 #!/usr/bin/python 2 # 3 # Copyright (c) 2012 The Chromium OS Authors. All rights reserved. 4 # Use of this source code is governed by a BSD-style license that can be 5 # found in the LICENSE file. 6 7 """Unit tests for site_utils/base_event.py.""" 8 9 import logging 10 import mox 11 import unittest 12 13 # driver must be imported first due to circular imports in base_event and task 14 import driver # pylint: disable-msg=W0611 15 import base_event 16 import deduping_scheduler 17 import task 18 19 20 class FakeTask(task.Task): 21 """A mock Task that can optionally expect triggering.""" 22 def __init__(self, suite, build, pool, pymox): 23 super(FakeTask, self).__init__('Fake', suite, build, pool) 24 pymox.StubOutWithMock(self, 'Run') 25 26 27 def CanArm(self, scheduler): 28 scheduler.CheckHostsExist( 29 multiple_labels=mox.IgnoreArg()).AndReturn(['host1']) 30 31 32 def Arm(self): 33 """Expect to be triggered along with any other FakeTasks.""" 34 self.Run(mox.IgnoreArg(), 35 mox.IgnoreArg(), 36 mox.IgnoreArg(), 37 mox.IgnoreArg(), 38 mox.IgnoreArg()).InAnyOrder('tasks').AndReturn(True) 39 40 41 class FakeOneShot(FakeTask): 42 """A mock OneShotEvent that can be optionally set to expect triggering.""" 43 def __init__(self, suite, build, pool, pymox): 44 super(FakeOneShot, self).__init__(suite, build, pool, pymox) 45 46 47 def Arm(self): 48 """Expect to be triggered once, and to ask for self-destruction.""" 49 self.Run(mox.IgnoreArg(), 50 mox.IgnoreArg(), 51 mox.IgnoreArg(), 52 mox.IgnoreArg(), 53 mox.IgnoreArg()).AndReturn(False) 54 55 56 class BaseEventTest(mox.MoxTestBase): 57 """Unit tests for BaseEvent. 58 59 @var _TASKS: Specs for several tasks to run. 60 """ 61 62 63 _TASKS = [('suite1', 'build1', 'pool'), 64 ('suite2', 'build2', None), 65 ('suite2', 'build3', None), 66 ('suite3', 'build2', None)] 67 68 69 def setUp(self): 70 super(BaseEventTest, self).setUp() 71 self.sched = self.mox.CreateMock(deduping_scheduler.DedupingScheduler) 72 73 74 def CreateEvent(self): 75 """Return an instance of base_event.BaseEvent.""" 76 return base_event.BaseEvent('event', None, False) 77 78 79 def testEventDeduping(self): 80 """Tests that tasks are de-duped at BaseEvent creation.""" 81 tasks = [FakeTask(*task, pymox=self.mox) for task in self._TASKS] 82 tasks.append(FakeTask(*self._TASKS[0], pymox=self.mox)) 83 self.mox.ReplayAll() 84 85 event = self.CreateEvent() 86 event.tasks = tasks 87 self.assertEquals(len(event.tasks), len(self._TASKS)) 88 89 90 def testRecurringTasks(self): 91 """Tests that tasks are all run on Handle().""" 92 tasks = [FakeTask(*task, pymox=self.mox) for task in self._TASKS] 93 for task in tasks: 94 task.CanArm(self.sched) 95 task.Arm() 96 event = self.CreateEvent() 97 self.mox.ReplayAll() 98 99 event.tasks = tasks 100 event.Handle(self.sched, {}, "board1") 101 self.mox.VerifyAll() 102 103 # Ensure that all tasks are still around and can be Handle()'d again. 104 self.mox.ResetAll() 105 for task in tasks: 106 task.CanArm(self.sched) 107 task.Arm() 108 self.mox.ReplayAll() 109 event.Handle(self.sched, {}, "board1") 110 111 112 def testOneShotWithRecurringTasks(self): 113 """Tests that one-shot tasks are destroyed correctly.""" 114 tasks = [FakeTask(*task, pymox=self.mox) for task in self._TASKS] 115 all_tasks = tasks + [FakeOneShot(*self._TASKS[0], pymox=self.mox)] 116 for task in all_tasks: 117 task.CanArm(self.sched) 118 task.Arm() 119 event = self.CreateEvent() 120 self.mox.ReplayAll() 121 122 event.tasks = all_tasks 123 event.Handle(self.sched, {}, "board1") 124 self.mox.VerifyAll() 125 126 # Ensure that only recurring tasks can get Handle()'d again. 127 self.mox.ResetAll() 128 for task in tasks: 129 task.CanArm(self.sched) 130 task.Arm() 131 self.mox.ReplayAll() 132 event.Handle(self.sched, {}, "board1") 133 134 135 if __name__ == '__main__': 136 unittest.main() 137