machine_manager_unittest.py revision c454cee542ca459ef9bd87c9f72e81c822caf1e5
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 "", "", "") 36LABEL_MIX = label.MockLabel("mix", "image", CHROMEOS_ROOT, "mix", 37 ["daisy1", "daisy2", "lumpy3", "lumpy4"], 38 "", "", "") 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