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