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