From c7887a57de4b8a2edc910df55d88e2d8842c52ed Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Valim?= <jose.valim@dashbit.co>
Date: Fri, 7 May 2021 08:57:39 +0200
Subject: [PATCH] Remove sleep from tests

---
 lib/livebook/runtime/node_pool.ex        |  6 ++-
 test/livebook/runtime/node_pool_test.exs | 47 ++++++++++--------------
 2 files changed, 24 insertions(+), 29 deletions(-)

diff --git a/lib/livebook/runtime/node_pool.ex b/lib/livebook/runtime/node_pool.ex
index 7305f2e01..0adcfd3a2 100644
--- a/lib/livebook/runtime/node_pool.ex
+++ b/lib/livebook/runtime/node_pool.ex
@@ -59,7 +59,11 @@ defmodule Livebook.Runtime.NodePool do
 
   @impl GenServer
   def handle_info({:nodedown, node, _info}, state) do
-    _ = Process.send_after(self(), {:add_node, node}, state.buffer_time)
+    case state.buffer_time do
+      0 -> send(self(), {:add_node, node})
+      t -> Process.send_after(self(), {:add_node, node}, t)
+    end
+
     {:noreply, state}
   end
 
diff --git a/test/livebook/runtime/node_pool_test.exs b/test/livebook/runtime/node_pool_test.exs
index 1a61e41de..8c8cd1be0 100644
--- a/test/livebook/runtime/node_pool_test.exs
+++ b/test/livebook/runtime/node_pool_test.exs
@@ -3,13 +3,9 @@ defmodule Livebook.Runtime.NodePoolTest do
 
   alias Livebook.Runtime.NodePool
 
-  # Tests for Livebook.Runtime.NodePool
-  #
-  # Note:
-  #
-  #   We do not spawn actual nodes as it can be time
-  #   intensive (on low spec machines) and is generally
-  #   complicated.
+  # Note we do not spawn actual nodes as it can be time
+  # intensive (on low spec machines) and is generally
+  # complicated.
 
   describe "start_link" do
     test "correctly starts a registered GenServer", config do
@@ -24,21 +20,16 @@ defmodule Livebook.Runtime.NodePoolTest do
     test "creates a new node name if pool is empty", config do
       start_supervised!({NodePool, name: config.test})
 
-      # Assert that we get a result and that it is an atom
       result = NodePool.get_name(config.test, node())
-      assert result
       assert is_atom(result)
+      assert result |> Atom.to_string() |> String.ends_with?(Atom.to_string(node()))
     end
 
     test "returns an existing name if pool is not empty", config do
       start_supervised!({NodePool, name: config.test, buffer_time: 0})
 
       name = NodePool.get_name(config.test, node())
-      send(config.test, {:nodedown, name, {}})
-
-      # Since we want the `:add_node` message  processed first
-      # before we call `get_name`, we wait
-      Process.sleep(1)
+      nodedown(config.test, name)
 
       assert NodePool.get_name(config.test, node()) == name
     end
@@ -47,11 +38,7 @@ defmodule Livebook.Runtime.NodePoolTest do
       start_supervised!({NodePool, name: config.test, buffer_time: 0})
 
       name = NodePool.get_name(config.test, node())
-      send(config.test, {:nodedown, name, {}})
-
-      # Since we want the `:add_node` message  processed first
-      # before we call `get_name`, we wait
-      Process.sleep(1)
+      nodedown(config.test, name)
 
       name = NodePool.get_name(config.test, node())
       assert NodePool.get_name(config.test, node()) != name
@@ -61,16 +48,20 @@ defmodule Livebook.Runtime.NodePoolTest do
   describe "on nodedown" do
     test "does not add node name to pool if not in generated_names", config do
       start_supervised!({NodePool, name: config.test, buffer_time: 0})
-
-      # Mock a nodedown
-      send(config.test, {:nodedown, :some_foo, {}})
-
-      # Since we want the `:add_node` message  processed first
-      # before we call `get_name`, we wait
-      Process.sleep(1)
-
-      # Verify that name is not in pool, by calling get_name/2
+      nodedown(config.test, :some_foo)
       assert NodePool.get_name(config.test, node()) != :some_foo
     end
   end
+
+  # Emulate node down and make sure it is processed
+  defp nodedown(process, node) do
+    send(process, {:nodedown, node, {}})
+
+    # Make sure the send was processed
+    _ = :sys.get_status(process)
+    # Make sure the send after message processed
+    _ = :sys.get_status(process)
+
+    :ok
+  end
 end