machine_manager_unittest.py revision f81680c018729fd4499e1e200d04b48c4b90127c 
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)
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      assert cm.machine_checksum, ("Could not find checksum for machine %s" %
27                                   machine_name)
28      self._all_machines.append(cm)
29
30CHROMEOS_ROOT = "/tmp/chromeos-root"
31MACHINE_NAMES = ["lumpy1", "lumpy2", "lumpy3", "daisy1", "daisy2"]
32LABEL_LUMPY = label.MockLabel("lumpy", "image", CHROMEOS_ROOT, "lumpy",
33                              ["lumpy1", "lumpy2", "lumpy3", "lumpy4"],
34                              "", "", "")
35LABEL_MIX = label.MockLabel("mix", "image", CHROMEOS_ROOT, "mix",
36                            ["daisy1", "daisy2", "lumpy3", "lumpy4"],
37                            "", "", "")
38
39
40class MachineManagerTest(unittest.TestCase):
41
42  def testAreAllMachineSame(self):
43    manager = MyMachineManager(CHROMEOS_ROOT)
44    for m in MACHINE_NAMES:
45      manager.AddMachine(m)
46    self.assertEqual(manager.AreAllMachineSame(LABEL_LUMPY), True)
47    self.assertEqual(manager.AreAllMachineSame(LABEL_MIX), False)
48
49  def testGetMachines(self):
50    manager = MyMachineManager(CHROMEOS_ROOT)
51    for m in MACHINE_NAMES:
52      manager.AddMachine(m)
53    names = [m.name for m in manager.GetMachines(LABEL_LUMPY)]
54    self.assertEqual(names, ["lumpy1", "lumpy2", "lumpy3"])
55
56  def testGetAvailableMachines(self):
57    manager = MyMachineManager(CHROMEOS_ROOT)
58    for m in MACHINE_NAMES:
59      manager.AddMachine(m)
60    for m in manager._all_machines:
61      if int(m.name[-1]) % 2:
62        manager._TryToLockMachine(m)
63    names = [m.name for m in manager.GetAvailableMachines(LABEL_LUMPY)]
64    self.assertEqual(names, ["lumpy1", "lumpy3"])
65
66if __name__ == "__main__":
67  unittest.main()
68