• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
9import mox
10import unittest
11
12# driver must be imported first due to circular imports in base_event and task
13import driver  # pylint: disable-msg=W0611
14import base_event
15import deduping_scheduler
16import task
17
18
19class 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
45class 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
61class 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
140if __name__ == '__main__':
141    unittest.main()
142