nvimdotfiles/lua/myworkflow/go-tests.lua

local M = {
    results = {}
}

local ns = vim.api.nvim_create_namespace("live-tests")

local test_function_query_string = [[
(
 (function_declaration
  name: (identifier) @name
  parameters:
    (parameter_list
     (parameter_declaration
      name: (identifier)
      type: (pointer_type
          (qualified_type
           package: (package_identifier) @_package_name
           name: (type_identifier) @_type_name)))))

 (#eq? @_package_name "testing")
 (#eq? @_type_name "T")
 (#eq? @name "%s")
)
]]

local t_run_query_string = [[
(

)
]]

local find_test_line = function(go_bufnr, name)
    local formatted = string.format(test_function_query_string, name)
    local query = vim.treesitter.query.parse("go", formatted)
    local parser = vim.treesitter.get_parser(go_bufnr, "go", {})
    local tree = parser:parse()[1]
    local root = tree:root()

    for id, node in query:iter_captures(root, go_bufnr, 0, -1) do
        if id == 1 then
            local range = { node:range() }
            return range[1]
        end
    end
end

local function clear(table)
    for k in pairs(table) do
        table[k] = nil
    end
end

vim.api.nvim_create_user_command('GoTest', function()
    M.goRunTests()
end, {})

vim.api.nvim_create_user_command('GoPrintResults', function()
    print(vim.inspect(M.results))
end, {})

vim.api.nvim_create_user_command('GoSetTestMarks', function()
    M.goSetAllMarks()
end, {})

vim.api.nvim_create_user_command('GoClearResults', function()
    clear(M.results)
end, {})


vim.api.nvim_create_user_command('GoClearTestMarks', function()
    local currbuf = vim.api.nvim_get_current_buf()
    vim.api.nvim_buf_clear_namespace(currbuf, ns, 0, -1)
end, {})

local errored = false
function M.goRunTests()
    -- TODO: replace jobstart/jobwait with system()
    local id = vim.fn.jobstart({ "go", "test", "./...", "-json" }, {
        stdout_buffered = true,
        on_stdout = function(_, data)
            if errored == true then
                errored = false
                return
            end

            if not data then
                return
            end

            clear(M.results)

            for _, line in ipairs(data) do
                if line == "" then goto continue end
                local decoded = vim.json.decode(line)

                if decoded.Test == nil then
                    goto continue
                end
                if decoded.Action == 'pass' or decoded.Action == 'fail' then
                    M.results[decoded.Test] = { Action = decoded.Action }
                end

                ::continue::
            end
        end,
        on_stderr = function(_, data)
            for _, value in pairs(data) do
                if value and value ~= "" then
                    errored = true
                    print('Error while running tests: ', value)
                end
            end
        end,
        on_exit = function()
            M.goSetAllMarks()
        end
    })
end

function M.goSetAllMarks()
    local files = vim.fn.system({ "find", "-type", "f", "-name", "*.go" })

    local sep = "\n"
    local filetable = {}
    for str in string.gmatch(files, "([^" .. sep .. "]+)") do
        table.insert(filetable, str)
    end

    for _, file in ipairs(filetable) do
        vim.fn.bufadd(file)
        vim.fn.bufload(file)
        local bufno = vim.fn.bufnr(file)
        M.goSetMarks(bufno)
    end
end

function M.goSetMarks(currbuf)
    vim.api.nvim_buf_clear_namespace(currbuf, ns, 0, -1)
    vim.diagnostic.set(ns, currbuf, {}, {})
    local text = {}
    local failed = {}

    for test, val in pairs(M.results) do
        local linenum = find_test_line(currbuf, test)
        if linenum == nil then goto continue end

        if val.Action == 'pass' then
            text = { "✔" }
            vim.api.nvim_buf_set_extmark(currbuf, ns, linenum, 0, {
                virt_text = { text },
            })
        elseif val.Action == 'fail' then
            table.insert(failed, {
                bufnr = currbuf,
                lnum = linenum,
                col = 0,
                severity = vim.diagnostic.severity.ERROR,
                source = "go-test",
                message = "Test Failed",
                user_data = {},
            })
        end

        ::continue::
    end
    vim.diagnostic.set(ns, currbuf, failed, {})
end

return M