Improve slave-preference function

scripts/master/unittests/master_utils_test.py

 #!/usr/bin/env python
 # Copyright 2013 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Source file for master_utils testcases."""
 
 
 import unittest
 
 import test_env  # pylint: disable=W0611,W0403
 
 from buildbot.process.properties import Properties
 from master import master_utils
 
 
+# Use this instead of random.choice on function return for determinism
+# and to check the full choice range.
+def mock_random_choice(iterable):
+  return list(iterable)
+
+
 class MasterUtilsTest(unittest.TestCase):
 
   def testPartition(self):
     partitions = master_utils.Partition([(1, 'a'),
                                          (2, 'b'),
                                          (3, 'c'),
                                          ], 2)
     self.assertEquals([['a', 'b'], ['c']], partitions)
 
 
@@ skipping 12 matching lines @@
     self.name = name
     self.slave = MockSlave(name, properties)
 
 class PreferredBuilderNextSlaveFuncTest(unittest.TestCase):
   def testNextSlave(self):
     builder1 = MockBuilder('builder1')
     builder2 = MockBuilder('builder2')
     builder3 = MockBuilder('builder3')
 
     slaves = [
-        MockSlaveBuilder('slave1', {'preferred_builder': 'builder1'}),
-        MockSlaveBuilder('slave2', {'preferred_builder': 'builder2'}),
-        MockSlaveBuilder('slave3', {'preferred_builder': 'builder3'}),
+        MockSlaveBuilder('s1', {'preferred_builder': 'builder1'}),
+        MockSlaveBuilder('s2', {'preferred_builder': 'builder2'}),
+        MockSlaveBuilder('s3', {'preferred_builder': 'builder3'}),
+        MockSlaveBuilder('s4', {'preferred_builder': 'builder1'}),
+        MockSlaveBuilder('s5', {'preferred_builder': 'builder2'}),
+        MockSlaveBuilder('s6', {'preferred_builder': 'builder3'}),
+        # Fall-over pool with no preference.
+        MockSlaveBuilder('s7', {'preferred_builder': None}),
+        MockSlaveBuilder('s8', {'preferred_builder': None}),
     ]
 
-    f = master_utils.PreferredBuilderNextSlaveFunc()
-    self.assertEqual('slave1', f(builder1, slaves).name)
-    self.assertEqual('slave2', f(builder2, slaves).name)
-    self.assertEqual('slave3', f(builder3, slaves).name)
+    f = lambda builder, slaves: (
+        set([s.name for s in master_utils.PreferredBuilderNextSlaveFunc(
+            mock_random_choice)(builder, slaves)]))
+    self.assertEqual(set(['s1', 's4']), f(builder1, slaves))
+    self.assertEqual(set(['s2', 's5']), f(builder2, slaves))
+    self.assertEqual(set(['s3', 's6']), f(builder3, slaves))
 
     # Remove slave 3.
     del slaves[2]
 
-    # When there is no slave that matches preferred_builder,
-    # any slave builder might be chosen.
-    self.assertTrue(f(builder3, slaves).name in ['slave1', 'slave2'])
+    # There's still a preferred slave left.
+    self.assertEqual(set(['s6']), f(builder3, slaves))
+
+    # Remove slave 6.
+    del slaves[4]

[tandrii(chromium), 2016/10/21 10:39:17]

def remove(slaves):
  iterate and find and remove

remove('s6') is nicer.

[Michael Achenbach, 2016/10/21 13:17:43]

Done.

+
+    # No preferred slave. Slave will be choosen from fall-over pool (i.e.
+    # slaves with no preference).
+    self.assertEqual(set(['s7', 's8']), f(builder3, slaves))
+
+    # Remove slaves 7 and 8.
+    del slaves[4]
+    del slaves[4]
+
+    # When the fall-over pool is depleted, slaves will be chosen from slave
+    # groups with most capacity (here: all the remaining ones are equal).
+    self.assertEqual(set(['s1', 's2', 's4', 's5']), f(builder3, slaves))
+
+    # Remove slave 1:
+    del slaves[0]
+
+    # Now only slaves preferring builder2 have most capacity.
+    self.assertEqual(set(['s2', 's5']), f(builder3, slaves))
 
   def testNextSlaveEmpty(self):
     builder = MockBuilder('builder')
     slaves = []
 
-    f = master_utils.PreferredBuilderNextSlaveFunc()
+    f = master_utils.PreferredBuilderNextSlaveFunc(mock_random_choice)
 
     self.assertIsNone(f(builder, slaves))
 
 
 if __name__ == '__main__':
   unittest.main()