Fix #24 (#25)

* Add test to repro the issue.
* Prevent undo from leaving behind added end text.
* Use --headless in tests to prevent UI open-close flickering when
  running many tests.

lua/nvim-treesitter/endwise.lua

@@ -95,8 +95,8 @@ local function add_end_node(indent_node_range, endable_node_range, end_text, shi
 
     local cursor_indentation = indentation .. string.rep(tabstr(), shiftcount)
 
-    vim.fn.append(crow, indentation .. end_text .. trailing_end_text)
     vim.fn.setline(crow, cursor_indentation .. trailing_cursor_text)
+    vim.fn.append(crow, indentation .. end_text .. trailing_end_text)
     vim.fn.cursor(crow, #cursor_indentation + 1)
 end
 

tests/endwise/undo.rb

@@ -0,0 +1,45 @@
+config({
+  extension: "rb",
+  command: 'ExecuteCRTwiceAndUndo',
+  overrides: <<~INIT_LUA
+vim.opt.expandtab = true
+vim.opt.shiftwidth = 2
+INIT_LUA
+})
+
+test "ruby, undo, clean undo-block after multiple <CR>s before end of file", <<~END
+-def foo█
+-
++def foo
++
+END
+
+test "ruby, undo, clean undo-block after multiple <CR>s at end of file", <<~END
+-
+-def foo█
++
++def foo
+END
+
+config({
+  extension: "lua",
+  command: 'ExecuteCRTwiceAndUndo',
+  overrides: <<~INIT_LUA
+vim.opt.expandtab = true
+vim.opt.shiftwidth = 2
+INIT_LUA
+})
+
+test "lua, undo, clean undo-block after multiple <CR>s before end of file", <<~END
+-function foo()█
+-
++function foo()
++
+END
+
+test "lua, undo, clean undo-block after multiple <CR>s at end of file", <<~END
+-
+-function foo()█
++
++function foo()
+END

tests/init.lua

@@ -1,6 +1,7 @@
 vim.opt.runtimepath:append('.')
 vim.opt.runtimepath:append('../nvim-treesitter')
 vim.opt.runtimepath:append('../playground')
+vim.opt.showmode = false
 require('nvim-treesitter.configs').setup {
     playground = {
         enable = true,
@@ -10,10 +11,38 @@ require('nvim-treesitter.configs').setup {
     },
 }
 
-function ExecuteCR(n)
-    vim.schedule_wrap(function()
-    local keys = vim.api.nvim_replace_termcodes(string.rep('l', n)..'a<CR>', true, false, true)
-    vim.cmd([[ autocmd User PostNvimTreesitterEndwiseCR lua vim.cmd('wq') ]])
+local function feedkeys(input)
+    local keys = vim.api.nvim_replace_termcodes(input, true, false, true)
     vim.fn.feedkeys(keys, 'n')
-    end)()
+end
+
+function ExecuteCR(n)
+    vim.schedule(function ()
+        vim.api.nvim_create_autocmd('User', {
+            pattern = 'PostNvimTreesitterEndwiseCR',
+            command = "silent wq"
+        })
+        feedkeys(string.rep('l', n)..'a<CR>')
+    end)
+end
+
+function ExecuteCRTwiceAndUndo(n)
+    local call_count = 0
+    local post_endwise_cb = function()
+        call_count = call_count + 1
+
+        if call_count < 2 then
+            feedkeys('<CR>')
+        else
+            vim.cmd([[ silent undo | silent wq ]])
+        end
+    end
+
+    vim.schedule(function ()
+        vim.api.nvim_create_autocmd('User', {
+            pattern = 'PostNvimTreesitterEndwiseCR',
+            callback = post_endwise_cb
+        })
+        feedkeys(string.rep('l', n)..'a<CR>')
+    end)
 end

tests/runner.rb

@@ -8,6 +8,7 @@ def config(opts)
   @config = {
     extension: "rb",
     overrides: "",
+    command: "ExecuteCR"
   }.merge(opts)
 end
 
@@ -37,15 +38,18 @@ def test(description, testcase)
     crow, ccol = get_cursor_pos(input)
     input = input.gsub(CURSOR, "")
     File.write(input_fname, input)
-    system("nvim", "-u", BASE_INIT_LUA, "+#{crow+1}", "-S", overrides, input_fname, "-c", "lua ExecuteCR(#{ccol-1})")
+    command = @config[:command]
+    system("nvim", "--headless", "-u", BASE_INIT_LUA, "+#{crow+1}", "-S", overrides, input_fname, "-c", "lua #{command}(#{ccol-1})")
     got = File.read(input_fname)
     if got != expected
+      puts ""
       puts "\e[31mFailed\e[0m: #{description}"
       puts "\e[34mInput\e[0m:", input.gsub(/\t/, "<tab>")
       puts "\e[34mGot\e[0m:", got.gsub(/\t/, "<tab>")
       puts "\e[34mExpected:\e[0m", expected.gsub(/\t/, "<tab>")
+      puts "\e[31m======\e[0m"
     else
-      puts "\e[32mSuccess\e[0m"
+      print "\e[32m.\e[0m"
     end
   end
 end
@@ -54,3 +58,5 @@ def test(description, testcase)
   @config = nil
   eval(File.read(fname))
 end
+
+puts ""