Mock RingBuilder.rebalance when testing ringbuilder cli warnings

Use mock to force explicit RingBuilder rebalance results so that the
test is focussed on the ringbuilder.py rebalance command behavior
when balance is (or is not) achieved. Avoids assumptions about
RingBuilder behavior.

Change-Id: I242ffc2f1a4f7b69a679832a65790223642dcea8
Closes-Bug: #1499015

test/unit/cli/test_ringbuilder.py

@@ -134,7 +134,7 @@ class TestCommands(unittest.TestCase, RunSwiftRingBuilderMixin):
                     msg += '%3d: %s\n' % (i, line)
                 self.fail(msg)
 
-    def create_sample_ring(self, part_power=6):
+    def create_sample_ring(self, part_power=6, overload=None):
         """
         Create a sample ring with four devices
 
@@ -151,6 +151,8 @@ class TestCommands(unittest.TestCase, RunSwiftRingBuilderMixin):
             pass
 
         ring = RingBuilder(part_power, 3, 1)
+        if overload is not None:
+            ring.set_overload(overload)
         ring.add_dev({'weight': 100.0,
                       'region': 0,
                       'zone': 0,
@@ -2095,36 +2097,51 @@ class TestCommands(unittest.TestCase, RunSwiftRingBuilderMixin):
         self.assertIsNone(ringbuilder.main(argv))
 
     def test_warn_at_risk(self):
-        # when the number of total part replicas (3 * 2 ** 4 = 48 in
-        # this ring) is less than the total units of weight (310 in this
-        # ring) the relative number of parts per unit of weight (called
-        # weight_of_one_part) is less than 1 - and each whole part
-        # placed takes up a larger ratio of the fractional number of
-        # parts the device wants - so it's much more difficult to
-        # satisfy a device's weight exactly - that is to say less parts
-        # to go around tends to make things lumpy
-        self.create_sample_ring(4)
-        ring = RingBuilder.load(self.tmpfile)
-        ring.devs[0]['weight'] = 10
-        ring.save(self.tmpfile)
+        # check that warning is generated when rebalance does not achieve
+        # satisfactory balance
+        self.create_sample_ring()
+        orig_rebalance = RingBuilder.rebalance
+        fake_balance = 6
+
+        def fake_rebalance(builder_instance, *args, **kwargs):
+            parts, balance, removed_devs = orig_rebalance(builder_instance)
+            return parts, fake_balance, removed_devs
+
         argv = ["", self.tmpfile, "rebalance"]
-        self.assertSystemExit(EXIT_WARNING, ringbuilder.main, argv)
+        with mock.patch("swift.common.ring.builder.RingBuilder.rebalance",
+                        fake_rebalance):
+            self.assertSystemExit(EXIT_WARNING, ringbuilder.main, argv)
+
+        # even when some overload is allowed
+        self.create_sample_ring(overload=0.05)
+        argv = ["", self.tmpfile, "rebalance"]
+        with mock.patch("swift.common.ring.builder.RingBuilder.rebalance",
+                        fake_rebalance):
+            self.assertSystemExit(EXIT_WARNING, ringbuilder.main, argv)
 
     def test_no_warn_when_balanced(self):
-        # when the number of total part replicas (3 * 2 ** 10 = 3072 in
-        # this ring) is larger than the total units of weight (310 in
-        # this ring) the relative number of parts per unit of weight
-        # (called weight_of_one_part) is more than 1 - and each whole
-        # part placed takes up a smaller ratio of the larger number of
-        # parts the device wants - so it's much easier to satisfy a
-        # device's weight exactly - that is to say more parts to go
-        # around tends to smooth things out
-        self.create_sample_ring(10)
-        ring = RingBuilder.load(self.tmpfile)
-        ring.devs[0]['weight'] = 10
-        ring.save(self.tmpfile)
+        # check that no warning is generated when satisfactory balance is
+        # achieved...
+        self.create_sample_ring()
+        orig_rebalance = RingBuilder.rebalance
+        fake_balance = 5
+
+        def fake_rebalance(builder_instance, *args, **kwargs):
+            parts, balance, removed_devs = orig_rebalance(builder_instance)
+            return parts, fake_balance, removed_devs
+
         argv = ["", self.tmpfile, "rebalance"]
-        self.assertSystemExit(EXIT_SUCCESS, ringbuilder.main, argv)
+        with mock.patch("swift.common.ring.builder.RingBuilder.rebalance",
+                        fake_rebalance):
+            self.assertSystemExit(EXIT_SUCCESS, ringbuilder.main, argv)
+
+        # ...or balance is within permitted overload
+        self.create_sample_ring(overload=0.06)
+        fake_balance = 6
+        argv = ["", self.tmpfile, "rebalance"]
+        with mock.patch("swift.common.ring.builder.RingBuilder.rebalance",
+                        fake_rebalance):
+            self.assertSystemExit(EXIT_SUCCESS, ringbuilder.main, argv)
 
     def test_invalid_device_name(self):
         self.create_sample_ring()