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