feat: re-generate testify lookup (#128)

README.md

@@ -367,23 +367,21 @@ more information on this.
 
 ### Testify suites
 
-> [!WARNING] 
+> [!WARNING]
 > This feature comes with some caveats and nuances, which is why it
 > is not enabled by default. I advise you to only enable this if you need it.
 
 There are some real shenaningans going on behind the scenes to make this work.
-😅 First, a lookup of "receiver type-to-suite test function" will be created of
-all Go test files in your project. Then, the generated Neotest node tree is
-modified by mutating private attributes and merging of nodes to avoid
+😅 First, an in-memory lookup of "receiver type-to-suite test function" will be
+created of all Go test files in your project. Then, the generated Neotest node
+tree is modified by mutating private attributes and merging of nodes to avoid
 duplicates. I'm personally a bit afraid of the maintenance burden of this
 feature... 🙈
 
-> [!NOTE] 
-> Right now, there is no way to update the lookup other than restarting
-> Neotest/Neovim. So in case you are implementing a new suite, please restart to
-> see the new suites/tests appear in e.g. the summary window. Also, nested tests
-> or table tests are not supported. All of this can be remedied at any time.
-> Feel free to dig in and open a PR!
+> [!NOTE]
+> Right now, nested tests and table tests are not supported. All of this
+> can be remedied at any time by extending the treesitter queries. Feel free to
+> dig in and open a PR!
 
 ## 🙏 PRs are welcome
 

lua/neotest-golang/features/testify/lookup.lua

@@ -1,6 +1,5 @@
 --- Lookup table for renaming Neotest namespaces (receiver type to testify suite function).
 
-local lib = require("neotest-golang.lib")
 local query = require("neotest-golang.features.testify.query")
 
 local M = {}
@@ -63,55 +62,66 @@ M.query = [[
             (identifier))))))
 ]]
 
---- The lookup table.
---- @type table<string, table>
-local lookup_table = {}
+local function create_lookup_manager()
+  local lookup_table = {}
 
---- Get the current lookup table, generating it if empty.
---- @return table<string, table> The lookup table containing testify suite information
-function M.get()
-  if vim.tbl_isempty(lookup_table) then
-    lookup_table = M.generate()
-  end
-  return lookup_table
+  return {
+    init = function(file_paths)
+      for _, file_path in ipairs(file_paths) do
+        lookup_table[file_path] = M.generate_data(file_path)
+      end
+      return lookup_table
+    end,
+    create = function(file_path)
+      if not lookup_table[file_path] then
+        lookup_table[file_path] = M.generate_data(file_path)
+      end
+      return lookup_table
+    end,
+    get = function()
+      return lookup_table
+    end,
+    clear = function()
+      lookup_table = {}
+    end,
+  }
 end
 
---- Generate the lookup table for testify suites.
+-- Create an instance of the lookup manager
+local lookup_manager = create_lookup_manager()
+
+--- Public lookup functions.
+M.initialize_lookup = lookup_manager.init
+M.create_lookup = lookup_manager.create
+M.get_lookup = lookup_manager.get
+M.clear_lookup = lookup_manager.clear
+
+--- Generate the lookup data for the given file.
 --- @return table<string, table> The generated lookup table
-function M.generate()
-  local cwd = vim.fn.getcwd()
-  local filepaths = lib.find.go_test_filepaths(cwd)
-  local lookup = {}
-  -- local global_suites = {}
+function M.generate_data(file_path)
+  local data = {}
 
   -- First pass: collect all data for the lookup table.
-  for _, filepath in ipairs(filepaths) do
-    local matches = query.run_query_on_file(filepath, M.query)
-
-    local package_name = matches.package
-        and matches.package[1]
-        and matches.package[1].text
-      or "unknown"
-
-    lookup[filepath] = {
-      package = package_name,
-      replacements = {},
-    }
-
-    for i, struct in ipairs(matches.suite_struct or {}) do
-      local func = matches.test_function[i]
-      if func then
-        lookup[filepath].replacements[struct.text] = func.text
-      end
+  local matches = query.run_query_on_file(file_path, M.query)
+
+  local package_name = matches.package
+      and matches.package[1]
+      and matches.package[1].text
+    or "unknown"
+
+  data = {
+    package = package_name,
+    replacements = {},
+  }
+
+  for i, struct in ipairs(matches.suite_struct or {}) do
+    local func = matches.test_function[i]
+    if func then
+      data.replacements[struct.text] = func.text
     end
   end
 
-  return lookup
-end
-
---- Clear the lookup table.
-function M.clear()
-  lookup_table = {}
+  return data
 end
 
 return M

lua/neotest-golang/features/testify/tree_modification.lua

@@ -1,33 +1,57 @@
 --- Functions to modify the Neotest tree, for testify suite support.
 
+local options = require("neotest-golang.options")
+local lib = require("neotest-golang.lib")
 local lookup = require("neotest-golang.features.testify.lookup")
 
 local M = {}
 
+local lookup_table = lookup.get_lookup()
+local ignore_filepaths_during_init = {}
+
 --- Modify the neotest tree, so that testify suites can be executed
 --- as Neotest namespaces.
 ---
 --- When testify tests are discovered, they are discovered with the Go receiver
 --- type as the Neotest namespace. However, to produce a valid test path,
 --- this receiver type must be replaced with the testify suite name in the
 --- Neotest tree.
+--- @param file_path string The path to the test file
 --- @param tree neotest.Tree The original neotest tree
 --- @return neotest.Tree The modified tree.
-function M.modify_neotest_tree(tree)
-  local lookup_map = lookup.get()
+function M.modify_neotest_tree(file_path, tree)
+  if vim.tbl_isempty(lookup_table) then
+    ignore_filepaths_during_init = lib.find.go_test_filepaths(vim.fn.getcwd())
+    lookup_table = lookup.initialize_lookup(ignore_filepaths_during_init)
+  end
+
+  if vim.tbl_contains(ignore_filepaths_during_init, file_path) then
+    -- some optimization;
+    -- ignore the first call, as it is handled by the initialization above.
+    for i, path in ipairs(ignore_filepaths_during_init) do
+      if path == file_path then
+        table.remove(ignore_filepaths_during_init, i)
+        break
+      end
+    end
+  else
+    -- after initialization, always update the lookup for the given filepath.
+    lookup_table = lookup.create_lookup(file_path)
+  end
 
-  if not lookup_map then
+  if not lookup_table then
     vim.notify(
       "No lookup found. Could not modify Neotest tree for testify suite support",
       vim.log.levels.WARN
     )
     return tree
   end
 
-  local modified_tree = M.replace_receiver_with_suite(tree:root(), lookup_map)
-  local tree_with_merged_namespaces =
-    M.merge_duplicate_namespaces(modified_tree)
-  return tree_with_merged_namespaces
+  local modified_tree = {}
+  modified_tree = M.replace_receiver_with_suite(tree:root(), lookup_table)
+  modified_tree = M.merge_duplicate_namespaces(modified_tree)
+
+  return modified_tree
 end
 
 --- Replace receiver methods with their corresponding test suites in the tree.

lua/neotest-golang/init.lua

@@ -14,11 +14,7 @@ local M = {}
 --- @field name string
 M.Adapter = {
   name = "neotest-golang",
-  init = function()
-    if options.get().testify_enabled == true then
-      testify.lookup.generate()
-    end
-  end,
+  init = function() end,
 }
 
 --- Find the project root directory given a current directory to work from.

lua/neotest-golang/query.lua

@@ -138,7 +138,7 @@ function M.detect_tests(file_path)
   local tree = lib.treesitter.parse_positions(file_path, query, opts)
 
   if options.get().testify_enabled == true then
-    tree = testify.tree_modification.modify_neotest_tree(tree)
+    tree = testify.tree_modification.modify_neotest_tree(file_path, tree)
   end
 
   return tree

tests/go/testify/lookup_spec.lua

@@ -1,13 +1,15 @@
 local _ = require("plenary")
 
 local options = require("neotest-golang.options")
+local lib = require("neotest-golang.lib")
 local testify = require("neotest-golang.features.testify")
 
 describe("Lookup", function()
   it("Generates tree replacement instructions", function()
     -- Arrange
     options.set({ testify_enabled = true }) -- enable testify
     local folderpath = vim.loop.cwd() .. "/tests/go"
+    local filepaths = lib.find.go_test_filepaths(vim.loop.cwd())
     local expected_lookup = {
       [folderpath .. "/positions_test.go"] = {
         package = "main",
@@ -33,10 +35,10 @@ describe("Lookup", function()
     }
 
     -- Act
-    testify.lookup.generate() -- generate lookup
+    testify.lookup.initialize_lookup(filepaths) -- generate lookup
 
     -- Assert
-    local lookup = testify.lookup.get()
+    local lookup = testify.lookup.get_lookup()
     assert.are.same(vim.inspect(expected_lookup), vim.inspect(lookup))
     assert.are.same(expected_lookup, lookup)
   end)

tests/go/testify/positions_spec.lua

@@ -3,6 +3,7 @@ local _ = require("plenary")
 
 local adapter = require("neotest-golang")
 local options = require("neotest-golang.options")
+local lib = require("neotest-golang.lib")
 local testify = require("neotest-golang.features.testify")
 
 local function compareIgnoringKeys(t1, t2, ignoreKeys)
@@ -81,8 +82,8 @@ describe("With testify_enabled=true", function()
     local test_filepath = vim.loop.cwd()
       .. "/tests/go/testify/positions_test.go"
     options.set({ testify_enabled = true }) -- enable testify
-    testify.lookup.generate() -- generate lookup
-
+    local filepaths = lib.find.go_test_filepaths(test_filepath)
+    testify.lookup.initialize_lookup(filepaths) -- generate lookup
     local expected = {
       {
         id = test_filepath,