From 4e813784cd1a9700bea96ed8c39720c0674ddb87 Mon Sep 17 00:00:00 2001
From: Lukasz Samson <lukaszsamson@gmail.com>
Date: Wed, 27 Dec 2023 18:17:19 +0100
Subject: [PATCH] register context variables in ex_unit setup and setup_all

---
 lib/elixir_sense/core/metadata_builder.ex     | 56 ++++++++++++++++---
 .../core/metadata_builder_test.exs            | 37 ++++++++++++
 2 files changed, 85 insertions(+), 8 deletions(-)

diff --git a/lib/elixir_sense/core/metadata_builder.ex b/lib/elixir_sense/core/metadata_builder.ex
index 389fe674..ba1a94c9 100644
--- a/lib/elixir_sense/core/metadata_builder.ex
+++ b/lib/elixir_sense/core/metadata_builder.ex
@@ -632,9 +632,9 @@ defmodule ElixirSense.Core.MetadataBuilder do
   # ex_unit describe
   defp pre(
          {:describe, meta, [name, _body]} = ast,
-         state
+         state = %{scopes: [[atom | _] | _]}
        )
-       when is_binary(name) do
+       when is_binary(name) and is_atom(atom) do
     line = Keyword.fetch!(meta, :line)
     column = Keyword.fetch!(meta, :column)
 
@@ -649,9 +649,9 @@ defmodule ElixirSense.Core.MetadataBuilder do
   # ex_unit not implemented test
   defp pre(
          {:test, meta, [name]},
-         state
+         state = %{scopes: [[atom | _] | _]}
        )
-       when is_binary(name) do
+       when is_binary(name) and is_atom(atom) do
     def_name = ex_unit_test_name(state, name)
     line = Keyword.fetch!(meta, :line)
     column = Keyword.fetch!(meta, :column)
@@ -668,9 +668,9 @@ defmodule ElixirSense.Core.MetadataBuilder do
   # ex_unit test without context
   defp pre(
          {:test, meta, [name, body]},
-         state
+         state = %{scopes: [[atom | _] | _]}
        )
-       when is_binary(name) do
+       when is_binary(name) and is_atom(atom) do
     def_name = ex_unit_test_name(state, name)
     line = Keyword.fetch!(meta, :line)
     column = Keyword.fetch!(meta, :column)
@@ -687,9 +687,9 @@ defmodule ElixirSense.Core.MetadataBuilder do
   # ex_unit test with context
   defp pre(
          {:test, meta, [name, param, body]},
-         state
+         state = %{scopes: [[atom | _] | _]}
        )
-       when is_binary(name) do
+       when is_binary(name) and is_atom(atom) do
     def_name = ex_unit_test_name(state, name)
     line = Keyword.fetch!(meta, :line)
     column = Keyword.fetch!(meta, :column)
@@ -703,6 +703,46 @@ defmodule ElixirSense.Core.MetadataBuilder do
     pre_func(ast_without_params, state, meta, def_name, [param])
   end
 
+  # ex_unit setup with context
+  defp pre(
+         {setup, meta, [param, body]},
+         state = %{scopes: [[atom | _] | _]}
+       )
+       when setup in [:setup, :setup_all] and is_atom(atom) do
+    line = Keyword.fetch!(meta, :line)
+    column = Keyword.fetch!(meta, :column)
+
+    # NOTE this name is not 100% correct - ex_unit uses counters instead of line but it's too complicated
+    def_name = :"__ex_unit_#{setup}_#{line}"
+    ast_without_params = {:def, meta, [{def_name, add_no_call([]), []}, [], body]}
+
+    state =
+      state
+      |> add_call_to_line({nil, setup, 2}, {line, column})
+
+    pre_func(ast_without_params, state, meta, def_name, [param])
+  end
+
+  # ex_unit setup without context
+  defp pre(
+         {setup, meta, [body]},
+         state = %{scopes: [[atom | _] | _]}
+       )
+       when setup in [:setup, :setup_all] and is_atom(atom) do
+    line = Keyword.fetch!(meta, :line)
+    column = Keyword.fetch!(meta, :column)
+
+    # NOTE this name is not 100% correct - ex_unit uses counters instead of line but it's too complicated
+    def_name = :"__ex_unit_#{setup}_#{line}"
+    ast_without_params = {:def, meta, [{def_name, add_no_call([]), []}, [], body]}
+
+    state =
+      state
+      |> add_call_to_line({nil, setup, 2}, {line, column})
+
+    pre_func(ast_without_params, state, meta, def_name, [{:_, [line: line, column: column], nil}])
+  end
+
   # function head with guards
   defp pre(
          {def_name, meta, [{:when, _, [{name, meta2, params}, guards]}, body]},
diff --git a/test/elixir_sense/core/metadata_builder_test.exs b/test/elixir_sense/core/metadata_builder_test.exs
index cf89402f..6917c13d 100644
--- a/test/elixir_sense/core/metadata_builder_test.exs
+++ b/test/elixir_sense/core/metadata_builder_test.exs
@@ -646,6 +646,43 @@ defmodule ElixirSense.Core.MetadataBuilderTest do
            )
   end
 
+  test "variables are added to environment in ex_unit setup" do
+    state =
+      """
+      defmodule MyModuleTests do
+        use ExUnit.Case, async: true
+
+        setup_all %{some: some} do
+          IO.puts("")
+        end
+
+        setup %{some: other} do
+          IO.puts("")
+        end
+
+        setup do
+          IO.puts("")
+        end
+
+        setup :clean_up_tmp_directory
+
+        setup [:clean_up_tmp_directory, :another_setup]
+
+        setup {MyModule, :my_setup_function}
+      end
+      """
+      |> string_to_state
+
+    assert [%VarInfo{type: nil, scope_id: scope_id}] = state |> get_line_vars(5)
+    assert [%VarInfo{name: :some}] = state.vars_info_per_scope_id[scope_id]
+
+    assert [%VarInfo{type: nil, scope_id: scope_id}] = state |> get_line_vars(9)
+    assert [%VarInfo{name: :other}] = state.vars_info_per_scope_id[scope_id]
+
+    # we do not generate defs - ExUnit.Callbacks.__setup__ is too complicated and generates def names with counters, e.g.
+    # :"__ex_unit_setup_#{counter}_#{length(setup)}"
+  end
+
   test "variables from outside module are added to environment" do
     state =
       """