diff --git a/test/livebook_web/live/session_live_test.exs b/test/livebook_web/live/session_live_test.exs
index 3a275d96e..1fac408c4 100644
--- a/test/livebook_web/live/session_live_test.exs
+++ b/test/livebook_web/live/session_live_test.exs
@@ -119,7 +119,8 @@ defmodule LivebookWeb.SessionLiveTest do
       evaluate_setup(session.pid)
 
       section_id = insert_section(session.pid)
-      cell_id = insert_text_cell(session.pid, section_id, :code, "Process.sleep(50)")
+      {source, continue_fun} = source_for_blocking()
+      cell_id = insert_text_cell(session.pid, section_id, :code, source)
 
       {:ok, view, _} = live(conn, ~p"/sessions/#{session.id}")
 
@@ -129,6 +130,8 @@ defmodule LivebookWeb.SessionLiveTest do
 
       assert %{cell_infos: %{^cell_id => %{eval: %{status: :evaluating}}}} =
                Session.get_data(session.pid)
+
+      continue_fun.()
     end
 
     test "reevaluting the setup cell", %{conn: conn, session: session} do
@@ -1402,7 +1405,6 @@ defmodule LivebookWeb.SessionLiveTest do
 
       section_id = insert_section(session.pid)
       insert_cell_with_output(session.pid, section_id, {:text, "Hello from the app!"})
-      wait_for_session_update(session.pid)
 
       slug = Livebook.Utils.random_short_id()
 
@@ -1431,6 +1433,9 @@ defmodule LivebookWeb.SessionLiveTest do
         |> live(~p"/apps/#{slug}")
         |> follow_redirect(conn)
 
+      assert_receive {:app_updated,
+                      %{slug: ^slug, sessions: [%{app_status: %{execution: :executed}}]}}
+
       assert render(view) =~ "Hello from the app!"
 
       Livebook.App.close(app.pid)
diff --git a/test/support/test_helpers.ex b/test/support/test_helpers.ex
index be81be222..d892cb066 100644
--- a/test/support/test_helpers.ex
+++ b/test/support/test_helpers.ex
@@ -69,4 +69,29 @@ defmodule Livebook.TestHelpers do
     end
     |> Macro.to_string()
   end
+
+  @doc """
+  Builds code that awaits for a messages before finishing.
+
+  Returns `{code, continue_fun}`, where calling `continue_fun` should
+  continue execution. Embedded runtime must be used for this to work.
+  """
+  def source_for_blocking() do
+    name = Livebook.Utils.random_short_id() |> String.to_atom()
+
+    code =
+      quote do
+        # This test uses the Embedded runtime, so we can target the
+        # process by name
+        Process.register(self(), unquote(name))
+        receive do: (:stop -> :ok)
+      end
+      |> Macro.to_string()
+
+    ack_fun = fn ->
+      send(Process.whereis(name), :pong)
+    end
+
+    {code, ack_fun}
+  end
 end