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
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
206if __name__ == '__main__':
207    unittest.main()
208