More reliable drain and revive command assertions

(cherry picked from commit 6f17c6c)

Author: michaelklishin

deps/rabbitmq_cli/test/test_helper.exs

@@ -250,6 +250,23 @@ defmodule TestHelper do
     end
   end
 
+  def await_condition(fun, timeout) do
+    retries = Integer.floor_div(timeout, 50)
+    await_condition_with_retries(fun, retries)
+  end
+
+  def await_condition_with_retries(_fun, 0) do
+    throw({:error, "Condition did not materialize"})
+  end
+  def await_condition_with_retries(fun, retries_left) do
+    case fun.() do
+      true -> :ok
+      _    ->
+          :timer.sleep(50)
+          await_condition_with_retries(fun, retries_left - 1)
+    end
+  end
+
   def is_rabbitmq_app_running() do
     :rabbit_misc.rpc_call(get_rabbit_hostname(), :rabbit, :is_booted, [])
   end

deps/rabbitmq_cli/test/upgrade/drain_command_test.exs

@@ -50,8 +50,8 @@ defmodule DrainCommandTest do
   test "run: puts target node into maintenance mode", context do
     assert not is_draining_node()
     assert :ok == @command.run([], context[:opts])
-    assert is_draining_node()
 
+    await_condition(fn -> is_draining_node() end, 7000)
     revive_node()
   end
 end

deps/rabbitmq_cli/test/upgrade/revive_command_test.exs

@@ -50,8 +50,8 @@ defmodule ReviveCommandTest do
   test "run: puts target node into regular operating mode", context do
     assert not is_draining_node()
     drain_node()
-    assert is_draining_node()
+    await_condition(fn -> is_draining_node() end, 7000)
     assert :ok == @command.run([], context[:opts])
-    assert not is_draining_node()
+    await_condition(fn -> not is_draining_node() end, 7000)
   end
 end