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