Handle doctests reports when there are multiple assertions at once

lib/livebook/runtime/evaluator/doctests.ex

@@ -57,20 +57,32 @@ defmodule Livebook.Runtime.Evaluator.Doctests do
   end
 
   defp report_doctest_result(%{state: {:failed, failure}} = test, lines) do
-    line = test.tags.doctest_line
+    doctest_line = test.tags.doctest_line
-    [prompt_line | _] = lines = Enum.drop(lines, line - 1)
+    [prompt_line | _] = lines = Enum.drop(lines, doctest_line - 1)
-
-    interval =
-      lines
-      |> Enum.take_while(&(not end_of_doctest?(&1)))
-      |> length()
-      |> Kernel.-(1)
 
     # TODO: end_line must come from Elixir to be reliable
+    doctest_lines = Enum.take_while(lines, &(not end_of_doctest?(&1)))
+
+    interval =
+      with {:error, %ExUnit.AssertionError{}, [{_, _, _, location} | _]} <- failure,
+           assertion_line = location[:line],
+           true <- is_integer(assertion_line) and assertion_line >= doctest_line do
+        length(doctest_lines) -
+          length(
+            doctest_lines
+            |> Enum.drop(assertion_line - doctest_line)
+            |> Enum.drop_while(&prompt?(&1))
+            |> Enum.drop_while(&(not prompt?(&1)))
+          )
+      else
+        _ ->
+          length(doctest_lines)
+      end
+
     result = %{
       column: count_columns(prompt_line, 0),
-      line: line,
+      line: doctest_line,
-      end_line: interval + line,
+      end_line: interval + doctest_line - 1,
       state: :failed,
       contents: IO.iodata_to_binary(format_failure(failure, test))
     }
@@ -82,6 +94,16 @@ defmodule Livebook.Runtime.Evaluator.Doctests do
   defp count_columns("\t" <> rest, counter), do: count_columns(rest, counter + 2)
   defp count_columns(_, counter), do: counter
 
+  defp prompt?(line) do
+    case String.trim_leading(line) do
+      "iex>" <> _ -> true
+      "iex(" <> _ -> true
+      "...>" <> _ -> true
+      "...(" <> _ -> true
+      _ -> false
+    end
+  end
+
   defp end_of_doctest?(line) do
     case String.trim_leading(line) do
       "" -> true

test/livebook/runtime/evaluator_test.exs

@@ -397,34 +397,32 @@ defmodule Livebook.Runtime.EvaluatorTest do
 
       refute Code.ensure_loaded?(Livebook.Runtime.EvaluatorTest.Exited)
     end
+  end
 
-    @tag :with_ebin_path
+  describe "doctests" do
-    test "runs doctests when a module is defined", %{evaluator: evaluator} do
+    @describetag :with_ebin_path
+
+    test "assertions", %{evaluator: evaluator} do
       code = ~S'''
-      defmodule Livebook.Runtime.EvaluatorTest.Doctests do
+      defmodule Livebook.Runtime.EvaluatorTest.DoctestsAssertions do
         @moduledoc """
 
             iex> raise "oops"
             ** (ArgumentError) not oops
 
-            iex> 1 +
-            :who_knows
-
-            iex> 1 = 2
-
             iex> require ExUnit.Assertions
             ...> ExUnit.Assertions.assert false
         """
 
         @doc """
-            iex> Livebook.Runtime.EvaluatorTest.Doctests.data()
+            iex> Livebook.Runtime.EvaluatorTest.DoctestsAssertions.data()
             %{
               name: "Amy Santiago",
               description: "nypd detective",
               precinct: 99
             }
 
-          iex> Livebook.Runtime.EvaluatorTest.Doctests.data()
+          iex> Livebook.Runtime.EvaluatorTest.DoctestsAssertions.data()
           %{name: "Jake Peralta", description: "NYPD detective"}
         """
         def data() do
@@ -434,22 +432,6 @@ defmodule Livebook.Runtime.EvaluatorTest do
             precinct: 99
           }
         end
-
-        @doc """
-            iex> Livebook.Runtime.EvaluatorTest.Doctests.raise_with_stacktrace()
-            :what
-        """
-        def raise_with_stacktrace() do
-          Enum.map(1, & &1)
-        end
-
-        @doc """
-            iex> Livebook.Runtime.EvaluatorTest.Doctests.exit()
-            :what
-        """
-        def exit() do
-          Process.exit(self(), :shutdown)
-        end
       end
       '''
 
@@ -472,62 +454,110 @@ defmodule Livebook.Runtime.EvaluatorTest do
       assert_receive {:runtime_evaluation_output, :code_1,
                       {:doctest_result, %{line: 7, state: :evaluating}}}
 
+      assert_receive {:runtime_evaluation_output, :code_1,
+                      {
+                        :doctest_result,
+                        %{
+                          column: 6,
+                          contents: "\e[31mExpected truthy, got false\e[0m",
+                          end_line: 8,
+                          line: 7,
+                          state: :failed
+                        }
+                      }}
+
+      assert_receive {:runtime_evaluation_output, :code_1,
+                      {:doctest_result, %{line: 12, state: :evaluating}}}
+
+      assert_receive {:runtime_evaluation_output, :code_1,
+                      {:doctest_result, %{line: 12, state: :success}}}
+
+      assert_receive {:runtime_evaluation_output, :code_1,
+                      {:doctest_result, %{line: 19, state: :evaluating}}}
+
+      assert_receive {:runtime_evaluation_output, :code_1,
+                      {
+                        :doctest_result,
+                        %{column: 4, contents: _, end_line: 20, line: 19, state: :failed}
+                      }}
+    end
+
+    test "multiple assertions at once", %{evaluator: evaluator} do
+      code = ~S'''
+      defmodule Livebook.Runtime.EvaluatorTest.DoctestsMiddle do
+        @moduledoc """
+
+            iex> 1 + 1
+            2
+            iex> 1 + 2
+            :wrong
+            iex> 1 + 3
+            4
+
+        """
+      end
+      '''
+
+      Evaluator.evaluate_code(evaluator, code, :code_1, [])
+
+      assert_receive {:runtime_evaluation_output, :code_1,
+                      {:doctest_result, %{line: 4, state: :evaluating}}}
+
       assert_receive {:runtime_evaluation_output, :code_1,
                       {:doctest_result,
                        %{
                          column: 6,
-                         contents:
+                         contents: _,
-                           "\e[31mDoctest did not compile, got: (TokenMissingError) " <> _,
+                         end_line: 7,
-                         end_line: 8,
+                         line: 4,
-                         line: 7,
                          state: :failed
                        }}}
+    end
+
+    test "runtime errors", %{evaluator: evaluator} do
+      code = ~S'''
+      defmodule Livebook.Runtime.EvaluatorTest.DoctestsRuntime do
+        @moduledoc """
+
+            iex> 1 = 2
+
+        """
+
+        @doc """
+            iex> Livebook.Runtime.EvaluatorTest.DoctestsRuntime.raise_with_stacktrace()
+            :what
+        """
+        def raise_with_stacktrace() do
+          Enum.map(1, & &1)
+        end
+
+        @doc """
+            iex> Livebook.Runtime.EvaluatorTest.DoctestsRuntime.exit()
+            :what
+        """
+        def exit() do
+          Process.exit(self(), :shutdown)
+        end
+      end
+      '''
+
+      Evaluator.evaluate_code(evaluator, code, :code_1, [])
 
       assert_receive {:runtime_evaluation_output, :code_1,
-                      {:doctest_result, %{line: 10, state: :evaluating}}}
+                      {:doctest_result, %{line: 4, state: :evaluating}}}
 
       assert_receive {:runtime_evaluation_output, :code_1,
                       {:doctest_result,
                        %{
                          column: 6,
                          contents: "\e[31mmatch (=) failed" <> _,
-                         end_line: 10,
+                         end_line: 4,
-                         line: 10,
+                         line: 4,
                          state: :failed
                        }}}
 
       assert_receive {:runtime_evaluation_output, :code_1,
-                      {:doctest_result, %{line: 12, state: :evaluating}}}
+                      {:doctest_result, %{line: 9, state: :evaluating}}}
-
-      assert_receive {:runtime_evaluation_output, :code_1,
-                      {
-                        :doctest_result,
-                        %{
-                          column: 6,
-                          contents: "\e[31mExpected truthy, got false\e[0m",
-                          end_line: 13,
-                          line: 12,
-                          state: :failed
-                        }
-                      }}
-
-      assert_receive {:runtime_evaluation_output, :code_1,
-                      {:doctest_result, %{line: 17, state: :evaluating}}}
-
-      assert_receive {:runtime_evaluation_output, :code_1,
-                      {:doctest_result, %{line: 17, state: :success}}}
-
-      assert_receive {:runtime_evaluation_output, :code_1,
-                      {:doctest_result, %{line: 24, state: :evaluating}}}
-
-      assert_receive {:runtime_evaluation_output, :code_1,
-                      {
-                        :doctest_result,
-                        %{column: 4, contents: _, end_line: 25, line: 24, state: :failed}
-                      }}
-
-      assert_receive {:runtime_evaluation_output, :code_1,
-                      {:doctest_result, %{line: 36, state: :evaluating}}}
 
       assert_receive {:runtime_evaluation_output, :code_1,
                       {
@@ -537,14 +567,14 @@ defmodule Livebook.Runtime.EvaluatorTest do
                           contents:
                             "\e[31m** (Protocol.UndefinedError) protocol Enumerable not implemented for 1 of type Integer. " <>
                               _,
-                          end_line: 37,
+                          end_line: 10,
-                          line: 36,
+                          line: 9,
                           state: :failed
                         }
                       }}
 
       assert_receive {:runtime_evaluation_output, :code_1,
-                      {:doctest_result, %{line: 44, state: :evaluating}}}
+                      {:doctest_result, %{line: 17, state: :evaluating}}}
 
       assert_receive {:runtime_evaluation_output, :code_1,
                       {
@@ -552,12 +582,42 @@ defmodule Livebook.Runtime.EvaluatorTest do
                         %{
                           column: 6,
                           contents: "\e[31m** (EXIT from #PID<" <> _,
-                          end_line: 45,
+                          end_line: 18,
-                          line: 44,
+                          line: 17,
                           state: :failed
                         }
                       }}
     end
+
+    test "invalid", %{evaluator: evaluator} do
+      code = ~S'''
+      defmodule Livebook.Runtime.EvaluatorTest.DoctestsInvalid do
+        @doc """
+
+            iex> 1 +
+            :who_knows
+
+        """
+        def foo, do: :ok
+      end
+      '''
+
+      Evaluator.evaluate_code(evaluator, code, :code_1, [])
+
+      assert_receive {:runtime_evaluation_output, :code_1,
+                      {:doctest_result, %{line: 4, state: :evaluating}}}
+
+      assert_receive {:runtime_evaluation_output, :code_1,
+                      {:doctest_result,
+                       %{
+                         column: 6,
+                         contents:
+                           "\e[31mDoctest did not compile, got: (TokenMissingError) " <> _,
+                         end_line: 5,
+                         line: 4,
+                         state: :failed
+                       }}}
+    end
   end
 
   describe "evaluate_code/6 identifier tracking" do