fix(nes): rendering edge cases and tests (#9)

* tests: add NES UI tests

* fix(nes): edge case around deletion only

* feat(tests): add signin tests and screenshots

Add signin fixture, screenshot files, and integration test for the Copilot
LSP sign in modal. The test simulates the device flow with code verification.

* feat(nes): support add only edits

Implement trailing whitespace trimming to handle add only edits and update UI
calculations for proper row placement. Adjust tests, mocks, and screenshots to
reflect the new behavior.

Author: tris203

lua/copilot-lsp/handlers.lua

@@ -111,7 +111,7 @@ local function sign_in(client_id, bufnr)
     client:request(methods.signIn, vim.empty_dict(), nil, bufnr)
 end
 
----@param res {busy: boolean, kind: 'Normal'|'Error'|'Warning'|'Incative', message: string}
+---@param res {busy: boolean, kind: 'Normal'|'Error'|'Warning'|'Inactive', message: string}
 M["didChangeStatus"] = function(err, res, ctx)
     if err then
         return

lua/copilot-lsp/nes/ui.lua

@@ -23,12 +23,16 @@ function M.clear_suggestion(bufnr, ns_id)
     vim.b[bufnr].nes_state = nil
 end
 
+local function trim_end(s)
+    return s:gsub("%s+$", "")
+end
+
 ---@private
 ---@param suggestion copilotlsp.InlineEdit
 ---@return nes.LineCalculationResult
 function M._calculate_lines(suggestion)
     local deleted_lines_count = suggestion.range["end"].line - suggestion.range.start.line
-    local added_lines = vim.split(suggestion.newText, "\n")
+    local added_lines = vim.split(trim_end(suggestion.newText), "\n")
     local added_lines_count = suggestion.newText == "" and 0 or #added_lines
     local same_line = 0
 
@@ -38,25 +42,40 @@ function M._calculate_lines(suggestion)
         same_line = 1
     end
 
+    -- if
+    --     suggestion.range.start.line == suggestion.range["end"].line
+    --     and suggestion.range.start.character == suggestion.range["end"].character
+    -- then
+    --     --add only
+    --     TODO: Do we need to position specifically for add only?
+    --     UI tests seem to say no
+    -- end
+
     -- Calculate positions for delete highlight extmark
     ---@type nes.DeleteExtmark
     local delete_extmark = {
         row = suggestion.range.start.line,
-        end_row = suggestion.range["end"].line + 1,
+        end_row = (
+            suggestion.range["end"].character ~= 0 and suggestion.range["end"].line + 1
+            or suggestion.range["end"].line
+        ),
     }
 
     -- Calculate positions for virtual lines extmark
     ---@type nes.VirtLinesExtmark
     local virt_lines_extmark = {
-        row = suggestion.range["end"].line,
+        row = (
+            suggestion.range["end"].character ~= 0 and suggestion.range["end"].line
+            or suggestion.range["end"].line - 1
+        ),
         virt_lines_count = added_lines_count,
     }
 
     -- Calculate positions for floating window
     ---@type nes.FloatWin
     local float_win = {
         height = #added_lines,
-        row = suggestion.range["end"].line + deleted_lines_count + 1,
+        row = suggestion.range["end"].line + deleted_lines_count + (suggestion.range["end"].character ~= 0 and 1 or 0),
     }
 
     return {

tests/fixtures/addonly_edit.txt

@@ -0,0 +1,3 @@
+1 line
+2 line
+4 line

tests/fixtures/multiline_edit.txt

@@ -0,0 +1,3 @@
+line one
+line two
+line three

tests/fixtures/removal_edit.txt

@@ -0,0 +1,3 @@
+line one
+line two
+line three

tests/fixtures/sameline_edit.txt

@@ -0,0 +1,3 @@
+aaa
+bbb
+ccc

tests/fixtures/signin.txt

@@ -0,0 +1 @@
+Sign In Test

tests/mock_lsp.lua

@@ -0,0 +1,119 @@
+local M = {}
+
+M.messages = {}
+
+---@param td lsp.VersionedTextDocumentIdentifier
+---@return copilotlsp.copilotInlineEditResponse
+local function getNesResponse(td)
+    local filename = vim.fs.basename(vim.uri_to_fname(td.uri))
+    ---@type table<string, copilotlsp.copilotInlineEditResponse>
+    local responses = {
+        ["sameline_edit.txt"] = {
+            edits = {
+                {
+                    command = { title = "mock", command = "mock" },
+                    range = {
+                        start = { line = 0, character = 0 },
+                        ["end"] = { line = 0, character = 3 },
+                    },
+                    textDocument = td,
+                    text = "xyz",
+                    ---@diagnostic disable-next-line: assign-type-mismatch
+                    newText = nil,
+                },
+            },
+        },
+        ["multiline_edit.txt"] = {
+            edits = {
+                {
+                    command = { title = "mock", command = "mock" },
+                    range = {
+                        start = { line = 0, character = 0 },
+                        ["end"] = { line = 1, character = 8 },
+                    },
+                    textDocument = td,
+                    text = "new line one\nnew line two",
+                    ---@diagnostic disable-next-line: assign-type-mismatch
+                    newText = nil,
+                },
+            },
+        },
+        ["removal_edit.txt"] = {
+            edits = {
+                {
+                    command = { title = "mock", command = "mock" },
+                    range = {
+                        start = { line = 1, character = 0 },
+                        ["end"] = { line = 2, character = 0 },
+                    },
+                    textDocument = td,
+                    text = "",
+                    ---@diagnostic disable-next-line: assign-type-mismatch
+                    newText = nil,
+                },
+            },
+        },
+        ["addonly_edit.txt"] = {
+            edits = {
+                {
+                    command = { title = "mock", command = "mock" },
+                    range = {
+                        start = { line = 2, character = 0 },
+                        ["end"] = { line = 2, character = 0 },
+                    },
+                    textDocument = td,
+                    text = "line 3\n",
+                    ---@diagnostic disable-next-line: assign-type-mismatch
+                    newText = nil,
+                },
+            },
+        },
+    }
+    local response = responses[filename]
+    assert(response, "unhandled doc")
+    return response
+end
+
+function M.server()
+    local closing = false
+    local srv = {}
+
+    function srv.request(method, params, handler)
+        table.insert(M.messages, { method = method, params = params })
+        if method == "initialize" then
+            handler(nil, {
+                capabilities = {},
+            })
+        elseif method == "shutdown" then
+            handler(nil, nil)
+        elseif method == "textDocument/copilotInlineEdit" then
+            local response = getNesResponse(params.textDocument)
+            handler(nil, response)
+        else
+            assert(false, "Unhandled method: " .. method)
+        end
+    end
+
+    function srv.notify(method, params)
+        table.insert(M.messages, { method = method, params = params })
+        if method == "exit" then
+            closing = true
+        end
+    end
+
+    function srv.is_closing()
+        return closing
+    end
+
+    function srv.terminate()
+        closing = true
+    end
+
+    return srv
+end
+
+function M.Reset()
+    M.messages = {}
+end
+
+return M

tests/nes/test_nes.lua

@@ -0,0 +1,101 @@
+local eq = MiniTest.expect.equality
+local ref = MiniTest.expect.reference_screenshot
+
+local child = MiniTest.new_child_neovim()
+
+local T = MiniTest.new_set()
+T["nes"] = MiniTest.new_set({
+    hooks = {
+        pre_case = function()
+            child.restart({ "-u", "scripts/minimal_init.lua" })
+            child.lua_func(function()
+                vim.lsp.config("copilot", {
+                    cmd = require("tests.mock_lsp").server,
+                })
+                vim.lsp.enable("copilot")
+            end)
+        end,
+        post_once = child.stop,
+    },
+})
+
+T["nes"]["same line edit"] = function()
+    child.cmd("edit tests/fixtures/sameline_edit.txt")
+    ref(child.get_screenshot())
+    vim.uv.sleep(500)
+    local lsp_name = child.lua_func(function()
+        return vim.lsp.get_clients()[1].name
+    end)
+    eq(lsp_name, "copilot")
+    child.lua_func(function()
+        local copilot = vim.lsp.get_clients()[1]
+        require("copilot-lsp.nes").request_nes(copilot)
+    end)
+    vim.uv.sleep(500)
+    ref(child.get_screenshot())
+    child.lua_func(function()
+        require("copilot-lsp.nes").apply_pending_nes(0)
+    end)
+    ref(child.get_screenshot())
+end
+
+T["nes"]["multi line edit"] = function()
+    child.cmd("edit tests/fixtures/multiline_edit.txt")
+    ref(child.get_screenshot())
+    vim.uv.sleep(500)
+    local lsp_name = child.lua_func(function()
+        return vim.lsp.get_clients()[1].name
+    end)
+    eq(lsp_name, "copilot")
+    child.lua_func(function()
+        local copilot = vim.lsp.get_clients()[1]
+        require("copilot-lsp.nes").request_nes(copilot)
+    end)
+    vim.uv.sleep(500)
+    ref(child.get_screenshot())
+    child.lua_func(function()
+        require("copilot-lsp.nes").apply_pending_nes(0)
+    end)
+    ref(child.get_screenshot())
+end
+
+T["nes"]["removal edit"] = function()
+    child.cmd("edit tests/fixtures/removal_edit.txt")
+    ref(child.get_screenshot())
+    vim.uv.sleep(500)
+    local lsp_name = child.lua_func(function()
+        return vim.lsp.get_clients()[1].name
+    end)
+    eq(lsp_name, "copilot")
+    child.lua_func(function()
+        local copilot = vim.lsp.get_clients()[1]
+        require("copilot-lsp.nes").request_nes(copilot)
+    end)
+    vim.uv.sleep(500)
+    ref(child.get_screenshot())
+    child.lua_func(function()
+        require("copilot-lsp.nes").apply_pending_nes(0)
+    end)
+    ref(child.get_screenshot())
+end
+
+T["nes"]["add only edit"] = function()
+    child.cmd("edit tests/fixtures/addonly_edit.txt")
+    ref(child.get_screenshot())
+    vim.uv.sleep(500)
+    local lsp_name = child.lua_func(function()
+        return vim.lsp.get_clients()[1].name
+    end)
+    eq(lsp_name, "copilot")
+    child.lua_func(function()
+        local copilot = vim.lsp.get_clients()[1]
+        require("copilot-lsp.nes").request_nes(copilot)
+    end)
+    vim.uv.sleep(500)
+    ref(child.get_screenshot())
+    child.lua_func(function()
+        require("copilot-lsp.nes").apply_pending_nes(0)
+    end)
+    ref(child.get_screenshot())
+end
+return T