Fix tests

Author: frankroeder

tests/parrot/response_handler_spec.lua

@@ -2,23 +2,31 @@ local ResponseHandler = require("parrot.response_handler")
 local stub = require("luassert.stub")
 
 describe("ResponseHandler", function()
-  local mock_vim, mock_utils, mock_queries
+  local mock_utils, mock_queries
+  local original_vim_api = {}
 
   before_each(function()
-    mock_vim = {
-      api = {
-        nvim_get_current_buf = stub.new().returns(1),
-        nvim_get_current_win = stub.new().returns(1),
-        nvim_create_namespace = stub.new().returns(1),
-        nvim_buf_set_extmark = stub.new().returns(1),
-        nvim_buf_is_valid = stub.new().returns(true),
-        nvim_buf_get_extmark_by_id = stub.new().returns({ 1 }),
-        nvim_buf_set_lines = stub.new(),
-        nvim_win_get_cursor = stub.new().returns({ 1, 0 }),
-      },
-      split = stub.new().returns({ "test" }),
-      -- Remove vim.cmd to avoid potential issues with Vim options
-    }
+    -- Store original vim API functions
+    original_vim_api.nvim_get_current_buf = vim.api.nvim_get_current_buf
+    original_vim_api.nvim_get_current_win = vim.api.nvim_get_current_win
+    original_vim_api.nvim_create_namespace = vim.api.nvim_create_namespace
+    original_vim_api.nvim_buf_set_extmark = vim.api.nvim_buf_set_extmark
+    original_vim_api.nvim_buf_is_valid = vim.api.nvim_buf_is_valid
+    original_vim_api.nvim_buf_get_extmark_by_id = vim.api.nvim_buf_get_extmark_by_id
+    original_vim_api.nvim_buf_set_lines = vim.api.nvim_buf_set_lines
+    original_vim_api.nvim_win_get_cursor = vim.api.nvim_win_get_cursor
+    original_vim_api.vim_split = vim.split
+
+    -- Mock vim API functions
+    vim.api.nvim_get_current_buf = stub.new().returns(1)
+    vim.api.nvim_get_current_win = stub.new().returns(1)
+    vim.api.nvim_create_namespace = stub.new().returns(1)
+    vim.api.nvim_buf_set_extmark = stub.new().returns(1)
+    vim.api.nvim_buf_is_valid = stub.new().returns(true)
+    vim.api.nvim_buf_get_extmark_by_id = stub.new().returns({ 1 })
+    vim.api.nvim_buf_set_lines = stub.new()
+    vim.api.nvim_win_get_cursor = stub.new().returns({ 1, 0 })
+    vim.split = stub.new().returns({ "test" })
 
     mock_utils = {
       uuid = stub.new().returns("test-uuid"),
@@ -30,13 +38,21 @@ describe("ResponseHandler", function()
       get = stub.new().returns({ ns_id = 1, ex_id = 1 }),
     }
 
-    -- Use package.loaded instead of _G.vim
-    package.loaded.vim = mock_vim
     package.loaded["parrot.utils"] = mock_utils
   end)
 
   after_each(function()
-    package.loaded.vim = nil
+    -- Restore original vim API functions
+    vim.api.nvim_get_current_buf = original_vim_api.nvim_get_current_buf
+    vim.api.nvim_get_current_win = original_vim_api.nvim_get_current_win
+    vim.api.nvim_create_namespace = original_vim_api.nvim_create_namespace
+    vim.api.nvim_buf_set_extmark = original_vim_api.nvim_buf_set_extmark
+    vim.api.nvim_buf_is_valid = original_vim_api.nvim_buf_is_valid
+    vim.api.nvim_buf_get_extmark_by_id = original_vim_api.nvim_buf_get_extmark_by_id
+    vim.api.nvim_buf_set_lines = original_vim_api.nvim_buf_set_lines
+    vim.api.nvim_win_get_cursor = original_vim_api.nvim_win_get_cursor
+    vim.split = original_vim_api.vim_split
+
     package.loaded["parrot.utils"] = nil
   end)
 
@@ -73,7 +89,7 @@ describe("ResponseHandler", function()
   end)
 
   it("should not process if buffer is invalid", function()
-    mock_vim.api.nvim_buf_is_valid.returns(false)
+    vim.api.nvim_buf_is_valid.returns(false)
     local handler = ResponseHandler:new(mock_queries)
     handler:handle_chunk(1, nil)
     assert.are.same("", handler.response)