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/board_enumerator.py.""" 8 9 import mox, unittest 10 11 # driver must be imported first due to circular imports in base_event and task 12 import driver # pylint: disable-msg=W0611 13 import base_event, board_enumerator, build_event, deduping_scheduler 14 import forgiving_config_parser, manifest_versions, task, timed_event 15 16 import common 17 from autotest_lib.server import frontend 18 from constants import Labels 19 20 21 class DriverTest(mox.MoxTestBase): 22 """Unit tests for Driver.""" 23 24 _BOARDS = ['board1', 'board2'] 25 26 27 def setUp(self): 28 super(DriverTest, self).setUp() 29 self.afe = self.mox.CreateMock(frontend.AFE) 30 self.be = board_enumerator.BoardEnumerator(self.afe) 31 self.ds = deduping_scheduler.DedupingScheduler(self.afe) 32 self.mv = self.mox.CreateMock(manifest_versions.ManifestVersions) 33 34 self.config = forgiving_config_parser.ForgivingConfigParser() 35 36 self.nightly_bvt = task.Task(timed_event.Nightly.KEYWORD, '', '') 37 self.weekly_bvt = task.Task(timed_event.Weekly.KEYWORD, '', '') 38 self.new_build_bvt = task.Task(build_event.NewBuild.KEYWORD, '', '') 39 40 self.driver = driver.Driver(self.ds, self.be) 41 42 43 def _CreateMockEvent(self, klass): 44 event = self.mox.CreateMock(klass) 45 event.keyword = klass.KEYWORD 46 event.tasks = [] 47 return event 48 49 50 def _ExpectSetup(self): 51 mock_nightly = self._CreateMockEvent(timed_event.Nightly) 52 mock_weekly = self._CreateMockEvent(timed_event.Weekly) 53 mock_new_build = self._CreateMockEvent(build_event.NewBuild) 54 55 self.mox.StubOutWithMock(timed_event.Nightly, 'CreateFromConfig') 56 self.mox.StubOutWithMock(timed_event.Weekly, 'CreateFromConfig') 57 self.mox.StubOutWithMock(build_event.NewBuild, 'CreateFromConfig') 58 timed_event.Nightly.CreateFromConfig( 59 mox.IgnoreArg(), self.mv).AndReturn(mock_nightly) 60 timed_event.Weekly.CreateFromConfig( 61 mox.IgnoreArg(), self.mv).AndReturn(mock_weekly) 62 build_event.NewBuild.CreateFromConfig( 63 mox.IgnoreArg(), self.mv).AndReturn(mock_new_build) 64 return [mock_nightly, mock_weekly, mock_new_build] 65 66 67 def _ExpectTaskConfig(self): 68 self.config.add_section(timed_event.Nightly.KEYWORD) 69 self.config.add_section(timed_event.Weekly.KEYWORD) 70 self.mox.StubOutWithMock(task.Task, 'CreateFromConfigSection') 71 task.Task.CreateFromConfigSection( 72 self.config, timed_event.Nightly.KEYWORD).InAnyOrder().AndReturn( 73 (timed_event.Nightly.KEYWORD, self.nightly_bvt)) 74 task.Task.CreateFromConfigSection( 75 self.config, timed_event.Weekly.KEYWORD).InAnyOrder().AndReturn( 76 (timed_event.Weekly.KEYWORD, self.weekly_bvt)) 77 78 79 def _ExpectEnumeration(self): 80 """Expect one call to BoardEnumerator.Enumerate().""" 81 prefix = Labels.BOARD_PREFIX 82 mocks = [] 83 for board in self._BOARDS: 84 mocks.append(self.mox.CreateMock(frontend.Label)) 85 mocks[-1].name = prefix + board 86 self.afe.get_labels(name__startswith=prefix).AndReturn(mocks) 87 88 89 def _ExpectHandle(self, event, group): 90 """Make event report that it's handle-able, and expect it to be handle. 91 92 @param event: the mock event that expectations will be set on. 93 @param group: group to put new expectations in. 94 """ 95 bbs = {'branch': 'build-string'} 96 event.ShouldHandle().InAnyOrder(group).AndReturn(True) 97 for board in self._BOARDS: 98 event.GetBranchBuildsForBoard( 99 board).InAnyOrder(group).AndReturn(bbs) 100 event.Handle(mox.IgnoreArg(), bbs, board).InAnyOrder(group) 101 # Should happen once per loop, not once per Handle() 102 # http://crosbug.com/30642 103 event.UpdateCriteria().InAnyOrder(group) 104 105 106 def _ExpectNoHandle(self, event, group): 107 """Make event report that it's handle-able, but do not expect to 108 handle it. 109 110 @param event: the mock event that expectations will be set on. 111 @param group: group to put new expectations in. 112 """ 113 bbs = {'branch': 'build-string'} 114 event.ShouldHandle().InAnyOrder(group).AndReturn(True) 115 116 117 def testTasksFromConfig(self): 118 """Test that we can build a list of Tasks from a config.""" 119 self._ExpectTaskConfig() 120 self.mox.ReplayAll() 121 tasks = self.driver.TasksFromConfig(self.config) 122 self.assertTrue(self.nightly_bvt in tasks[timed_event.Nightly.KEYWORD]) 123 self.assertTrue(self.weekly_bvt in tasks[timed_event.Weekly.KEYWORD]) 124 125 126 def testTasksFromConfigRecall(self): 127 """Test that we can build a list of Tasks from a config twice.""" 128 events = self._ExpectSetup() 129 self._ExpectTaskConfig() 130 self.mox.ReplayAll() 131 132 self.driver.SetUpEventsAndTasks(self.config, self.mv) 133 for keyword, event in self.driver._events.iteritems(): 134 if keyword == timed_event.Nightly.KEYWORD: 135 self.assertTrue(self.nightly_bvt in event.tasks) 136 if keyword == timed_event.Weekly.KEYWORD: 137 self.assertTrue(self.weekly_bvt in event.tasks) 138 139 self.mox.UnsetStubs() 140 self.mox.VerifyAll() 141 142 self.mox.ResetAll() 143 self.config.remove_section(timed_event.Weekly.KEYWORD) 144 145 self._ExpectSetup() 146 self.mox.StubOutWithMock(task.Task, 'CreateFromConfigSection') 147 task.Task.CreateFromConfigSection( 148 self.config, timed_event.Nightly.KEYWORD).InAnyOrder().AndReturn( 149 (timed_event.Nightly.KEYWORD, self.nightly_bvt)) 150 self.mox.ReplayAll() 151 self.driver.SetUpEventsAndTasks(self.config, self.mv) 152 for keyword, event in self.driver._events.iteritems(): 153 if keyword == timed_event.Nightly.KEYWORD: 154 self.assertTrue(self.nightly_bvt in event.tasks) 155 elif keyword == timed_event.Weekly.KEYWORD: 156 self.assertFalse(self.weekly_bvt in event.tasks) 157 158 159 def testHandleAllEventsOnce(self): 160 """Test that all events being ready is handled correctly.""" 161 events = self._ExpectSetup() 162 self._ExpectEnumeration() 163 for event in events: 164 self._ExpectHandle(event, 'events') 165 self.mox.ReplayAll() 166 167 driver.POOL_SIZE = 1 168 self.driver.SetUpEventsAndTasks(self.config, self.mv) 169 self.driver.HandleEventsOnce(self.mv) 170 171 172 def testHandleNightlyEventOnce(self): 173 """Test that one ready event is handled correctly.""" 174 events = self._ExpectSetup() 175 self._ExpectEnumeration() 176 for event in events: 177 if event.keyword == timed_event.Nightly.KEYWORD: 178 self._ExpectHandle(event, 'events') 179 else: 180 event.ShouldHandle().InAnyOrder('events').AndReturn(False) 181 self.mox.ReplayAll() 182 183 driver.POOL_SIZE = 1 184 self.driver.SetUpEventsAndTasks(self.config, self.mv) 185 self.driver.HandleEventsOnce(self.mv) 186 187 188 def testForceOnceForBuild(self): 189 """Test that one event being forced is handled correctly.""" 190 events = self._ExpectSetup() 191 192 board = 'board' 193 type = 'release' 194 milestone = '00' 195 manifest = '200.0.02' 196 build = base_event.BuildName(board, type, milestone, manifest) 197 198 events[0].Handle(mox.IgnoreArg(), {milestone: [build]}, board, 199 force=True) 200 self.mox.ReplayAll() 201 202 self.driver.SetUpEventsAndTasks(self.config, self.mv) 203 self.driver.ForceEventsOnceForBuild([events[0].keyword], build) 204 205 206 if __name__ == '__main__': 207 unittest.main() 208
