machine_manager_unittest.py revision 2a370f7af34f6b126e80ef55891b9bff6a2858bb 
1#!/usr/bin/python
2
3# Copyright 2012 Google Inc. All Rights Reserved.
4
5"""Unittest for machine_manager."""
6import unittest
7
8import label
9import machine_manager
10
11
12class MyMachineManager(machine_manager.MachineManager):
13
14  def __init__(self, chromeos_root):
15    super(MyMachineManager, self).__init__(chromeos_root, 0, "average")
16
17  def _TryToLockMachine(self, cros_machine):
18    self._machines.append(cros_machine)
19    cros_machine.checksum = ""
20
21  def AddMachine(self, machine_name):
22    with self._lock:
23      for m in self._all_machines:
24        assert m.name != machine_name, "Tried to double-add %s" % machine_name
25      cm = machine_manager.MockCrosMachine(machine_name, self.chromeos_root,
26                                           "average")
27      assert cm.machine_checksum, ("Could not find checksum for machine %s" %
28                                   machine_name)
29      self._all_machines.append(cm)
30
31CHROMEOS_ROOT = "/tmp/chromeos-root"
32MACHINE_NAMES = ["lumpy1", "lumpy2", "lumpy3", "daisy1", "daisy2"]
33LABEL_LUMPY = label.MockLabel("lumpy", "image", CHROMEOS_ROOT, "lumpy",
34                              ["lumpy1", "lumpy2", "lumpy3", "lumpy4"],
35                              "", "", False, "")
36LABEL_MIX = label.MockLabel("mix", "image", CHROMEOS_ROOT, "mix",
37                            ["daisy1", "daisy2", "lumpy3", "lumpy4"],
38                            "", "", False, "")
39
40
41class MachineManagerTest(unittest.TestCase):
42
43  def testAreAllMachineSame(self):
44    manager = MyMachineManager(CHROMEOS_ROOT)
45    for m in MACHINE_NAMES:
46      manager.AddMachine(m)
47    self.assertEqual(manager.AreAllMachineSame(LABEL_LUMPY), True)
48    self.assertEqual(manager.AreAllMachineSame(LABEL_MIX), False)
49
50  def testGetMachines(self):
51    manager = MyMachineManager(CHROMEOS_ROOT)
52    for m in MACHINE_NAMES:
53      manager.AddMachine(m)
54    names = [m.name for m in manager.GetMachines(LABEL_LUMPY)]
55    self.assertEqual(names, ["lumpy1", "lumpy2", "lumpy3"])
56
57  def testGetAvailableMachines(self):
58    manager = MyMachineManager(CHROMEOS_ROOT)
59    for m in MACHINE_NAMES:
60      manager.AddMachine(m)
61    for m in manager._all_machines:
62      if int(m.name[-1]) % 2:
63        manager._TryToLockMachine(m)
64    names = [m.name for m in manager.GetAvailableMachines(LABEL_LUMPY)]
65    self.assertEqual(names, ["lumpy1", "lumpy3"])
66
67if __name__ == "__main__":
68  unittest.main()
69