diff --git a/lib/diff/tmp_dir.ex b/lib/diff/tmp_dir.ex
index 6b41db5..72db92f 100644
--- a/lib/diff/tmp_dir.ex
+++ b/lib/diff/tmp_dir.ex
@@ -32,33 +32,57 @@ defmodule Diff.TmpDir do
     dir
   end
 
+  def await_cleanup(pid) do
+    GenServer.call(__MODULE__, {:await_cleanup, pid}, 5000)
+  end
+
   defp track(path) do
     pid = self()
     :ets.insert(@table, {pid, path})
-    GenServer.cast(__MODULE__, {:monitor, pid})
+    GenServer.call(__MODULE__, {:monitor, pid})
   end
 
   @impl true
   def init(_opts) do
     Process.flag(:trap_exit, true)
     :ets.new(@table, [:named_table, :duplicate_bag, :public])
-    {:ok, %{monitors: MapSet.new()}}
+    {:ok, %{monitors: MapSet.new(), waiters: %{}}}
   end
 
   @impl true
-  def handle_cast({:monitor, pid}, state) do
+  def handle_call({:monitor, pid}, _from, state) do
     if pid in state.monitors do
-      {:noreply, state}
+      {:reply, :ok, state}
     else
       Process.monitor(pid)
-      {:noreply, %{state | monitors: MapSet.put(state.monitors, pid)}}
+      {:reply, :ok, %{state | monitors: MapSet.put(state.monitors, pid)}}
+    end
+  end
+
+  @impl true
+  def handle_call({:await_cleanup, pid}, from, state) do
+    if pid in state.monitors do
+      waiters = Map.update(state.waiters, pid, [from], &[from | &1])
+      {:noreply, %{state | waiters: waiters}}
+    else
+      {:reply, :ok, state}
     end
   end
 
   @impl true
   def handle_info({:DOWN, _ref, :process, pid, _reason}, state) do
     cleanup_pid(pid)
-    {:noreply, %{state | monitors: MapSet.delete(state.monitors, pid)}}
+
+    for from <- Map.get(state.waiters, pid, []) do
+      GenServer.reply(from, :ok)
+    end
+
+    {:noreply,
+     %{
+       state
+       | monitors: MapSet.delete(state.monitors, pid),
+         waiters: Map.delete(state.waiters, pid)
+     }}
   end
 
   @impl true
diff --git a/test/diff/tmp_dir_test.exs b/test/diff/tmp_dir_test.exs
index 4943135..7b10137 100644
--- a/test/diff/tmp_dir_test.exs
+++ b/test/diff/tmp_dir_test.exs
@@ -22,9 +22,8 @@ defmodule Diff.TmpDirTest do
         send(test_pid, {:paths, file, dir})
       end)
 
-    ref = Process.monitor(task_pid)
     assert_receive {:paths, file, dir}
-    wait_for_cleanup(task_pid, ref)
+    Diff.TmpDir.await_cleanup(task_pid)
 
     refute File.exists?(file)
     refute File.exists?(dir)
@@ -42,9 +41,8 @@ defmodule Diff.TmpDirTest do
         raise "crash"
       end)
 
-    ref = Process.monitor(task_pid)
     assert_receive {:paths, file, dir}
-    wait_for_cleanup(task_pid, ref)
+    Diff.TmpDir.await_cleanup(task_pid)
 
     refute File.exists?(file)
     refute File.exists?(dir)
@@ -65,9 +63,8 @@ defmodule Diff.TmpDirTest do
         send(test_pid, {:paths, paths})
       end)
 
-    ref = Process.monitor(task_pid)
     assert_receive {:paths, paths}
-    wait_for_cleanup(task_pid, ref)
+    Diff.TmpDir.await_cleanup(task_pid)
 
     for {file, dir} <- paths do
       refute File.exists?(file)
@@ -82,10 +79,4 @@ defmodule Diff.TmpDirTest do
     assert File.exists?(file)
     assert File.dir?(dir)
   end
-
-  defp wait_for_cleanup(task_pid, ref) do
-    assert_receive {:DOWN, ^ref, :process, ^task_pid, _}, 5000
-    # Sync with the GenServer to ensure the :DOWN cleanup has been processed
-    :sys.get_state(Diff.TmpDir)
-  end
 end
diff --git a/test/diff_web/integration_test.exs b/test/diff_web/integration_test.exs
index aac7e67..2a36b8c 100644
--- a/test/diff_web/integration_test.exs
+++ b/test/diff_web/integration_test.exs
@@ -31,6 +31,9 @@ defmodule DiffWeb.IntegrationTest do
         {:error, :not_found}
       end)
 
+      Diff.HexMock
+      |> stub(:diff, fn "phoenix", "1.4.5", "1.4.9" -> :error end)
+
       {:ok, _view, html} = live(conn, "/diff/phoenix/1.4.5..")
 
       # Should show generating state since we're resolving to latest version
@@ -46,6 +49,9 @@ defmodule DiffWeb.IntegrationTest do
         {:error, :not_found}
       end)
 
+      Diff.HexMock
+      |> stub(:diff, fn "nonexistent", "1.0.0", "2.0.0" -> :error end)
+
       {:ok, _view, html} = live(conn, "/diff/nonexistent/1.0.0..2.0.0")
 
       assert html =~ "Generating diffs"
diff --git a/test/diff_web/live/diff_live_view_test.exs b/test/diff_web/live/diff_live_view_test.exs
index 7e1ed1e..546a079 100644
--- a/test/diff_web/live/diff_live_view_test.exs
+++ b/test/diff_web/live/diff_live_view_test.exs
@@ -54,6 +54,9 @@ defmodule DiffWeb.DiffLiveViewTest do
         {:error, :not_found}
       end)
 
+      Diff.HexMock
+      |> stub(:diff, fn "phoenix", "1.4.5", "1.4.9" -> :error end)
+
       {:ok, _view, html} = live(conn, "/diff/phoenix/1.4.5..1.4.9")
 
       # Should show generating state when metadata is not found