Update floating builder logic, add to "chromiumos"

scripts/master/unittests/floating_builder_test.py

Index: scripts/master/unittests/floating_builder_test.py
diff --git a/scripts/master/unittests/floating_builder_test.py b/scripts/master/unittests/floating_builder_test.py
new file mode 100755
index 0000000000000000000000000000000000000000..b048be8b7f29d49f4daa97e4c0facf4dbb1e7a10
--- /dev/null
+++ b/scripts/master/unittests/floating_builder_test.py
@@ -0,0 +1,201 @@
+#!/usr/bin/env python
+# Copyright (c) 2012 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 floating builder testcases."""
+
+import calendar
+import datetime
+import itertools
+import os
+import time
+import unittest
+
+import test_env  # pylint: disable=W0611,W0403
+
+import mock
+
+from master import floating_builder as fb
+
+
+def _to_timestamp(dt):
+  # Calculate the offset between local timezone and UTC.
+  current_time = time.mktime(dt.timetuple())
+  offset = (datetime.datetime.fromtimestamp(current_time) -
+            datetime.datetime.utcfromtimestamp(current_time))
+
+  return calendar.timegm((dt - offset).timetuple())
+
+
+class _FakeSlaveStatus(object):
+  def __init__(self, name):
+    self.name = name
+    self.connect_times = []
+    self.last_message_received = None
+
+  def lastMessageReceived(self):
+    return self.last_message_received
+
+
+class _FakeSlave(object):
+  def __init__(self, slavename):
+    self.slavename = slavename
+    self.slave_status = None
+    self.offline = False
+
+  def _set_last_seen(self, now, **kwargs):
+    td = datetime.timedelta(**kwargs)
+    self.slave_status = _FakeSlaveStatus(self.slavename)
+    self.slave_status.last_message_received = _to_timestamp(now + td)
+
+  def __str__(self):
+    return self.slavename
+
+
+class _FakeBuilder(object):
+
+  def __init__(self, name, slaves):
+    self.name = name
+    self._all_slaves = slaves
+
+    self.botmaster = mock.MagicMock()
+    self.builder_status = mock.MagicMock()
+    self.builder_status.getSlaves.side_effect = lambda: [
+        s.slave_status for s in self._all_slaves
+        if s.slave_status]
+
+    self._online_slaves = ()
+    self._busy_slaves = ()
+
+  def __repr__(self):
+    return self.name
+
+  @property
+  def slaves(self):
+    return [_FakeSlaveBuilder(s, self)
+            for s in self._all_slaves
+            if s.slavename in self._online_slaves]
+
+  @property
+  def slavebuilders(self):
+    """Returns the list of slavebuilders that would be handed to
+    NextSlaveFunc.
+
+    This is the set of slaves that are available for scheduling. We derive
+    this by returning all slaves that are both online and not busy.
+    """
+    return self._get_slave_builders(lambda s:
+      s.slavename in self._online_slaves and
+      s.slavename not in self._busy_slaves)
+
+  def _get_slave_builders(self, fn):
+    return [_FakeSlaveBuilder(slave, self)
+            for slave in self._all_slaves
+            if fn(slave)]
+
+  def set_online_slaves(self, *slavenames):
+    self._online_slaves = set(slavenames)
+
+  def set_busy_slaves(self, *slavenames):
+    self._busy_slaves = set(slavenames)
+
+
+class _FakeSlaveBuilder(object):
+
+  def __init__(self, slave, builder):
+    self.slave = slave
+    self.builder = builder
+
+  def __repr__(self):
+    return '{%s/%s}' % (self.builder.name, self.slave.slavename)
+
+
+class FloatingBuilderTest(unittest.TestCase):
+
+  def setUp(self):
+    self._mocks = (
+      mock.patch('master.floating_builder._get_now'),
+      mock.patch('master.floating_builder.PokeBuilderTimer.reset'),
+    )
+    for patcher in self._mocks:
+      patcher.start()
+
+    # Mock current date/time.
+    self.now = datetime.datetime(2016, 1, 1, 8, 0, 0) # 1/1/2016 @8:00
+    fb._get_now.return_value = self.now
+
+    # Mock PokeBuilderTimer to record when the poke builder was set, but not
+    # actually schedule any reactor magic.
+    self.poke_delta = None
+    def record_poke_delta(delta):
+      self.poke_delta = delta
+    fb.PokeBuilderTimer.reset.side_effect = record_poke_delta
+
+    self._slaves = dict((s, _FakeSlave(s)) for s in (
+        'primary-a', 'primary-b', 'floating-a', 'floating-b',
+    ))
+
+    self.builder = _FakeBuilder(
+        'Test Builder',
+        [s[1] for s in sorted(self._slaves.iteritems())],
+    )
+
+  def tearDown(self):
+    for patcher in reversed(self._mocks):
+      patcher.stop()
+
+  def testPrimaryBuilderIsSelectedWhenAvailable(self):
+    fs = fb.FloatingSet()
+    fs.AddPrimary('primary-a')
+    fs.AddFloating('floating-a', 'floating-b')
+
+    self.builder.set_online_slaves('primary-a', 'floating-a', 'floating-b')
+
+    fnsf = fs.NextSlaveFunc(datetime.timedelta(seconds=10))
+    nsb = fnsf(self.builder, self.builder.slavebuilders)
+    self.assertIsNotNone(nsb)
+    self.assertEqual(nsb.slave.slavename, 'primary-a')
+
+  def testPrimaryBuilderIsSelectedWhenOneIsAvailableAndOneIsBusy(self):
+    fs = fb.FloatingSet()
+    fs.AddPrimary('primary-a', 'primary-b')
+    fs.AddFloating('floating-a', 'floating-b')
+
+    self.builder.set_online_slaves('primary-a', 'primary-b', 'floating-a',
+                                   'floating-b')
+    self.builder.set_busy_slaves('primary-a')
+
+    fnsf = fs.NextSlaveFunc(datetime.timedelta(seconds=10))
+    nsb = fnsf(self.builder, self.builder.slavebuilders)
+    self.assertIsNotNone(nsb)
+    self.assertEqual(nsb.slave.slavename, 'primary-b')
+
+  def testNoBuilderIsSelectedWhenPrimariesAreOfflineWithinGrace(self):
+    fs = fb.FloatingSet()
+    fs.AddPrimary('primary-a', 'primary-b')
+    fs.AddFloating('floating-a', 'floating-b')
+
+    self.builder.set_online_slaves('floating-a')
+    self._slaves['primary-b']._set_last_seen(self.now, seconds=-1)
+
+    fnsf = fs.NextSlaveFunc(datetime.timedelta(seconds=10))
+    nsb = fnsf(self.builder, self.builder.slavebuilders)
+    self.assertIsNone(nsb)
+    self.assertEqual(self.poke_delta, datetime.timedelta(seconds=9))
+
+  def testFloatingBuilderIsSelectedWhenPrimariesAreOfflineForAWhile(self):
+    fs = fb.FloatingSet()
+    fs.AddPrimary('primary-a', 'primary-b')
+    fs.AddFloating('floating-a', 'floating-b')
+
+    self.builder.set_online_slaves('floating-a')
+
+    fnsf = fs.NextSlaveFunc(datetime.timedelta(seconds=10))
+    nsb = fnsf(self.builder, self.builder.slavebuilders)
+    self.assertIsNotNone(nsb)
+    self.assertEqual(nsb.slave.slavename, 'floating-a')
+
+
+if __name__ == '__main__':
+  unittest.main()