From b335e1c9cb2b63cd038b5f5ca058f33d34893303 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Valim?= <jose.valim@dashbit.co>
Date: Wed, 23 Nov 2022 15:06:08 +0100
Subject: [PATCH] More doctests (#1544)

---
 lib/livebook/runtime/evaluator/doctests.ex | 57 ++++++++++++++++------
 test/livebook/runtime/evaluator_test.exs   | 16 +++++-
 2 files changed, 56 insertions(+), 17 deletions(-)

diff --git a/lib/livebook/runtime/evaluator/doctests.ex b/lib/livebook/runtime/evaluator/doctests.ex
index a516b6b73..9ef5f1d7b 100644
--- a/lib/livebook/runtime/evaluator/doctests.ex
+++ b/lib/livebook/runtime/evaluator/doctests.ex
@@ -31,7 +31,7 @@ if Version.compare(System.version(), "1.14.2") != :lt do
           delete_test_module(test_module)
 
         {:error, kind, error} ->
-          IO.warn(Exception.format(kind, error, []), [])
+          put_output({:error, colorize(:red, Exception.format(kind, error, [])), :other})
       end
 
       :ok
@@ -165,27 +165,50 @@ if Version.compare(System.version(), "1.14.2") != :lt do
 
       expected = diff[:right]
       got = diff[:left]
-
       source = String.trim(reason.doctest)
 
-      [
+      message_io =
+        if_io(reason.message != "Doctest failed", fn ->
+          message =
+            reason.message
+            |> String.replace_prefix("Doctest failed: ", "")
+            |> pad(@pad_size)
+
+          [colorize(:red, message), "\n"]
+        end)
+
+      source_io = [
         String.duplicate(" ", @pad_size),
         format_label("doctest"),
         "\n",
-        pad(source, @pad_size + 2),
-        "\n",
-        String.duplicate(" ", @pad_size),
-        format_label("expected"),
-        "\n",
-        String.duplicate(" ", @pad_size + 2),
-        expected,
-        "\n",
-        String.duplicate(" ", @pad_size),
-        format_label("got"),
-        "\n",
-        String.duplicate(" ", @pad_size + 2),
-        got
+        pad(source, @pad_size + 2)
       ]
+
+      expected_io =
+        if_io(expected, fn ->
+          [
+            "\n",
+            String.duplicate(" ", @pad_size),
+            format_label("expected"),
+            "\n",
+            String.duplicate(" ", @pad_size + 2),
+            expected
+          ]
+        end)
+
+      got_io =
+        if_io(got, fn ->
+          [
+            "\n",
+            String.duplicate(" ", @pad_size),
+            format_label("got"),
+            "\n",
+            String.duplicate(" ", @pad_size + 2),
+            got
+          ]
+        end)
+
+      message_io ++ source_io ++ expected_io ++ got_io
     end
 
     defp format_failure({kind, reason, stacktrace}, test) do
@@ -217,6 +240,8 @@ if Version.compare(System.version(), "1.14.2") != :lt do
       end
     end
 
+    defp if_io(value, fun), do: if(value, do: fun.(), else: [])
+
     defp format_stacktrace(stacktrace, test_case, test) do
       for entry <- stacktrace do
         message = format_stacktrace_entry(entry, test_case, test)
diff --git a/test/livebook/runtime/evaluator_test.exs b/test/livebook/runtime/evaluator_test.exs
index 74ada5194..bf67a3314 100644
--- a/test/livebook/runtime/evaluator_test.exs
+++ b/test/livebook/runtime/evaluator_test.exs
@@ -352,6 +352,15 @@ defmodule Livebook.Runtime.EvaluatorTest do
     test "runs doctests when a module is defined", %{evaluator: evaluator} do
       code = ~S'''
       defmodule Livebook.Runtime.EvaluatorTest.Doctests do
+        @moduledoc """
+
+            iex> raise "oops"
+            ** (ArgumentError) not oops
+
+            iex> 1 +
+            :who_knows
+        """
+
         @doc """
             iex> Livebook.Runtime.EvaluatorTest.Doctests.data()
             %{
@@ -392,7 +401,12 @@ defmodule Livebook.Runtime.EvaluatorTest do
       Evaluator.evaluate_code(evaluator, code, :code_1, [])
 
       assert_receive {:runtime_evaluation_output, :code_1, {:text, doctest_result}}
-      assert doctest_result =~ "4 doctests, 3 failures"
+
+      assert doctest_result =~ "6 doctests, 5 failures"
+      assert doctest_result =~ "Doctest did not compile, got: (TokenMissingError)"
+
+      assert doctest_result =~
+               "expected exception ArgumentError but got RuntimeError with message \"oops\""
 
       assert_receive {:runtime_evaluation_response, :code_1, {:ok, _}, metadata()}
     end