Synced repos to: 64740

1 files changed, 65 insertions, 0 deletions

diff --git a/crosperf/machine_manager_unittest.py b/crosperf/machine_manager_unittest.py
new file mode 100755
index 00000000..98baf456
--- /dev/null
+++ b/crosperf/machine_manager_unittest.py
@@ -0,0 +1,65 @@
+#!/usr/bin/python
+
+# Copyright 2012 Google Inc. All Rights Reserved.
+
+"""Unittest for machine_manager."""
+import unittest
+
+import label
+import machine_manager
+
+
+class MyMachineManager(machine_manager.MachineManager):
+
+  def __init__(self, chromeos_root):
+    super(MyMachineManager, self).__init__(chromeos_root)
+
+  def _TryToLockMachine(self, cros_machine):
+    self._machines.append(cros_machine)
+    cros_machine.checksum = ""
+
+  def AddMachine(self, machine_name):
+    with self._lock:
+      for m in self._all_machines:
+        assert m.name != machine_name, "Tried to double-add %s" % machine_name
+      cm = machine_manager.MockCrosMachine(machine_name, self.chromeos_root)
+      assert cm.machine_checksum, ("Could not find checksum for machine %s" %
+                                   machine_name)
+      self._all_machines.append(cm)
+
+CHROMEOS_ROOT = "/tmp/chromeos-root"
+MACHINE_NAMES = ["lumpy1", "lumpy2", "lumpy3", "daisy1", "daisy2"]
+LABEL_LUMPY = label.MockLabel("lumpy", "image", CHROMEOS_ROOT, "lumpy",
+                              ["lumpy1", "lumpy2", "lumpy3", "lumpy4"], "")
+LABEL_MIX = label.MockLabel("mix", "image", CHROMEOS_ROOT, "mix",
+                            ["daisy1", "daisy2", "lumpy3", "lumpy4"], "")
+
+
+class MachineManagerTest(unittest.TestCase):
+
+  def testAreAllMachineSame(self):
+    manager = MyMachineManager(CHROMEOS_ROOT)
+    for m in MACHINE_NAMES:
+      manager.AddMachine(m)
+    self.assertEqual(manager.AreAllMachineSame(LABEL_LUMPY), True)
+    self.assertEqual(manager.AreAllMachineSame(LABEL_MIX), False)
+
+  def testGetMachines(self):
+    manager = MyMachineManager(CHROMEOS_ROOT)
+    for m in MACHINE_NAMES:
+      manager.AddMachine(m)
+    names = [m.name for m in manager.GetMachines(LABEL_LUMPY)]
+    self.assertEqual(names, ["lumpy1", "lumpy2", "lumpy3"])
+
+  def testGetAvailableMachines(self):
+    manager = MyMachineManager(CHROMEOS_ROOT)
+    for m in MACHINE_NAMES:
+      manager.AddMachine(m)
+    for m in manager._all_machines:
+      if int(m.name[-1]) % 2:
+        manager._TryToLockMachine(m)
+    names = [m.name for m in manager.GetAvailableMachines(LABEL_LUMPY)]
+    self.assertEqual(names, ["lumpy1", "lumpy3"])
+
+if __name__ == "__main__":
+  unittest.main()