From 117685bdebcbd7a159b0c1abad3e6f344d3d899b Mon Sep 17 00:00:00 2001
From: Alexandre Pasmantier <alex.pasmant@gmail.com>
Date: Sun, 4 Aug 2024 00:29:32 +0200
Subject: [PATCH 1/4] tests: unit tests using luassert

---
 .gitignore                                    |   2 +
 lua/pymple/config.lua                         |  23 +-
 lua/pymple/init.lua                           |   1 +
 lua/pymple/resolve_imports.lua                |   8 +-
 lua/pymple/update_imports.lua                 |  25 +-
 lua/pymple/utils.lua                          |  78 ++++--
 .../contains_python_files/foo/bar/baz.py      |   0
 lua/tests/fixtures/utils/docstrings/misc_1.py |   4 +
 lua/tests/fixtures/utils/docstrings/misc_2.py |   4 +
 lua/tests/fixtures/utils/docstrings/misc_3.py |   6 +
 .../utils/docstrings/multiline_docstring.py   |   9 +
 .../fixtures/utils/docstrings/no_docstring.py |   4 +
 .../utils/docstrings/single_line_docstring.py |   6 +
 .../utils/no_python_files/bar/baz.txt         |   0
 {tests => lua/tests}/pymple_spec.lua          |   0
 lua/tests/resolve_imports_spec.lua            |  16 ++
 lua/tests/update_imports_spec.lua             |  62 +++++
 lua/tests/utils_spec.lua                      | 224 ++++++++++++++++++
 18 files changed, 434 insertions(+), 38 deletions(-)
 create mode 100644 lua/tests/fixtures/utils/contains_python_files/foo/bar/baz.py
 create mode 100644 lua/tests/fixtures/utils/docstrings/misc_1.py
 create mode 100644 lua/tests/fixtures/utils/docstrings/misc_2.py
 create mode 100644 lua/tests/fixtures/utils/docstrings/misc_3.py
 create mode 100644 lua/tests/fixtures/utils/docstrings/multiline_docstring.py
 create mode 100644 lua/tests/fixtures/utils/docstrings/no_docstring.py
 create mode 100644 lua/tests/fixtures/utils/docstrings/single_line_docstring.py
 create mode 100644 lua/tests/fixtures/utils/no_python_files/bar/baz.txt
 rename {tests => lua/tests}/pymple_spec.lua (100%)
 create mode 100644 lua/tests/resolve_imports_spec.lua
 create mode 100644 lua/tests/update_imports_spec.lua
 create mode 100644 lua/tests/utils_spec.lua

diff --git a/.gitignore b/.gitignore
index 6fd0a37..d122dd9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -39,3 +39,5 @@ luac.out
 *.x86_64
 *.hex
 
+# nvim help tags
+tags
diff --git a/lua/pymple/config.lua b/lua/pymple/config.lua
index 6720472..7910a7c 100644
--- a/lua/pymple/config.lua
+++ b/lua/pymple/config.lua
@@ -106,14 +106,31 @@ local default_logging_options = {
   level = "debug",
 }
 
----@alias Config { keymaps: Keymaps, create_user_commands: UserCommandOptions, update_imports: UpdateImportsOptions, logging: LoggingOptions}
+---@alias PythonOptions { virtual_env_names: string[] }
 
----@type Config
-M.default_config = {
+---@type PythonOptions
+local default_python_options = {
+  virtual_env_names = { ".venv" },
+}
+
+---@alias Config { keymaps: Keymaps, create_user_commands: UserCommandOptions, update_imports: UpdateImportsOptions, logging: LoggingOptions, python: PythonOptions}
+
+local default_config = {
   keymaps = default_keymaps,
   create_user_commands = default_user_command_options,
   update_imports = default_update_imports_options,
   logging = default_logging_options,
+  python = default_python_options,
 }
 
+---@type Config
+M.default_config = default_config
+
+---@param opts Config
+function M.set_config(opts)
+  M.config = vim.tbl_deep_extend("force", default_config, opts)
+end
+
+M.set_config(M.default_config)
+
 return M
diff --git a/lua/pymple/init.lua b/lua/pymple/init.lua
index 65d0626..bc0540f 100644
--- a/lua/pymple/init.lua
+++ b/lua/pymple/init.lua
@@ -26,6 +26,7 @@ local function setup(opts)
   opts = opts or {}
 
   setmetatable(opts, { __index = config.default_config })
+  config.set_config(opts)
 
   if opts.logging.enabled then
     opts.logging.enabled = nil
diff --git a/lua/pymple/resolve_imports.lua b/lua/pymple/resolve_imports.lua
index 7fd1343..c7609f5 100644
--- a/lua/pymple/resolve_imports.lua
+++ b/lua/pymple/resolve_imports.lua
@@ -19,15 +19,19 @@ local IMPORTABLE_SYMBOLS_PATTERNS = {
 
 ---@param args table
 ---@param symbol string
+---@param regexes string[]
 ---@return table
-local function add_symbol_regexes(args, symbol)
-  for _, pattern in ipairs(IMPORTABLE_SYMBOLS_PATTERNS) do
+local function add_symbol_regexes(args, symbol, regexes)
+  regexes = regexes or IMPORTABLE_SYMBOLS_PATTERNS
+  for _, pattern in ipairs(regexes) do
     table.insert(args, "-e")
     table.insert(args, string.format(pattern, symbol))
   end
   return args
 end
 
+M.add_symbol_regexes = add_symbol_regexes
+
 ---@param symbol string: the symbol for which to resolve an import
 ---@return string[] | nil: list of candidates
 function M.resolve_python_import(symbol)
diff --git a/lua/pymple/update_imports.lua b/lua/pymple/update_imports.lua
index e94eb69..8ef4dd7 100644
--- a/lua/pymple/update_imports.lua
+++ b/lua/pymple/update_imports.lua
@@ -20,6 +20,8 @@ local function build_filetypes_args(filetypes)
   return args
 end
 
+M.build_filetypes_args = build_filetypes_args
+
 --[[
 from path.to.dir import file
 from path.to.dir import (
@@ -31,7 +33,8 @@ from path.to.file import Something
 ---@param source string: The path to the source file/dir
 ---@param destination string: The path to the destination file/dir
 ---@param filetypes string[]: The filetypes to update imports for
-local function update_imports_split(source, destination, filetypes)
+local function update_imports_split(source, destination, filetypes, _job)
+  local __job = _job or jobs.gg_into_sed
   local cwd = vim.fn.getcwd()
 
   local source_relative = Path:new(source):make_relative(cwd)
@@ -71,10 +74,12 @@ local function update_imports_split(source, destination, filetypes)
       .. utils.escape_import_path(destination_module_name)
       .. "/'",
   }
-  jobs.gg_into_sed(gg_args_split, sed_args_base, true)
-  jobs.gg_into_sed(gg_args_split, sed_args_module, true)
+  __job(gg_args_split, sed_args_base, true)
+  __job(gg_args_split, sed_args_module, true)
 end
 
+M.update_imports_split = update_imports_split
+
 --[[
 from path.to.dir import file
 from path.to.dir import (
@@ -85,7 +90,8 @@ from path.to.dir import (
 ---@param source string: The path to the source file/dir
 ---@param destination string: The path to the destination file/dir
 ---@param filetypes string[]: The filetypes to update imports for
-local function update_imports_monolithic(source, destination, filetypes)
+local function update_imports_monolithic(source, destination, filetypes, _job)
+  local __job = _job or jobs.gg_into_sed
   local cwd = vim.fn.getcwd()
 
   -- path/to/here
@@ -102,21 +108,16 @@ local function update_imports_monolithic(source, destination, filetypes)
     string.format("'%s[\\.\\s]'", utils.escape_import_path(source_import_path)),
     ".",
   }
-  local sed_args = {
-    "'s/"
-      .. utils.escape_import_path(source_import_path)
-      .. "/"
-      .. utils.escape_import_path(destination_import_path)
-      .. "/'",
-  }
   local sed_args = string.format(
     "'s/%s\\([\\. ]\\)/%s\\1/'",
     utils.escape_import_path(source_import_path),
     utils.escape_import_path(destination_import_path)
   )
-  jobs.gg_into_sed(gg_args, { sed_args }, false)
+  __job(gg_args, { sed_args }, false)
 end
 
+M.update_imports_monolithic = update_imports_monolithic
+
 ---@param source string: The path to the source file/dir
 ---@param destination string: The path to the destination file/dir
 ---@param filetypes string[]: The filetypes to update imports for
diff --git a/lua/pymple/utils.lua b/lua/pymple/utils.lua
index b887849..85dc5f2 100644
--- a/lua/pymple/utils.lua
+++ b/lua/pymple/utils.lua
@@ -1,11 +1,11 @@
 M = {}
 
 local filetype = require("plenary.filetype")
-local config = require("pymple.config")
+local cfg = require("pymple.config")
 local log = require("pymple.log")
 
 ---@type number: The time to wait before refreshing open buffers
-DEFAULT_HANG_TIME = 1000
+local DEFAULT_HANG_TIME = 1000
 -- @param hang_time number: The time to wait before refreshing the buffers
 function M.async_refresh_buffers(hang_time)
   vim.defer_fn(function()
@@ -28,6 +28,7 @@ function M.check_binary_installed(binary_name)
   return 1 == vim.fn.executable(binary_name)
 end
 
+---Converts a path to an import path
 ---@param module_path string: The path to a python module
 ---@return string: The import path for the module
 function M.to_import_path(module_path)
@@ -35,18 +36,21 @@ function M.to_import_path(module_path)
   return result
 end
 
+---Splits an import path on the last separator
 ---@param import_path string: The import path to be split
----@return string, string: The base path and the last part of the import path
+---@return string | nil, string: The base path and the last part of the import path
 function M.split_import_on_last_separator(import_path)
   local base_path, module_name = import_path:match("(.-)%.?([^%.]+)$")
   return base_path, module_name
 end
 
+---Escapes a string to be used in a regex
 ---@param import_path string: The import path to be escaped
 function M.escape_import_path(import_path)
   return import_path:gsub("%.", [[\.]])
 end
 
+---Checks if a file is a python file
 ---@param path string: The path to the file
 ---@return boolean: Whether or not the file is a python file
 local function is_python_file(path)
@@ -55,6 +59,7 @@ end
 
 M.is_python_file = is_python_file
 
+---Recursively checks if a directory contains python files
 ---@param path string: The path to the directory
 ---@return boolean: Whether or not the directory contains python files
 local function recursive_dir_contains_python_files(path)
@@ -74,14 +79,14 @@ end
 
 M.recursive_dir_contains_python_files = recursive_dir_contains_python_files
 
----@param buf number: The buffer number
+---Finds the end line number of a docstring in a list of lines
+---@param lines string[]: The lines to search for a docstring
 ---@return number | nil: The height (in lines) of the docstring
-local function find_docstring_end_line_number(buf)
-  local lines = vim.api.nvim_buf_get_lines(buf, 0, -1, false)
+local function find_docstring_end_line_number_in_lines(lines)
   -- if the first line does not contain a docstring, return 0
   if not lines[1]:match('^"""') then
     return 0
-  elseif lines[1]:match('"""$') then
+  elseif lines[1]:match('""".*"""$') then
     return 1
   else
     for i = 2, #lines do
@@ -93,35 +98,66 @@ local function find_docstring_end_line_number(buf)
   return nil
 end
 
+M.find_docstring_end_line_number_in_lines =
+  find_docstring_end_line_number_in_lines
+
+---Finds the end line number of a file docstring
+---@param buf number: The buffer number
+---@return number | nil: The height (in lines) of the docstring
+local function find_docstring_end_line_number(buf)
+  local lines = vim.api.nvim_buf_get_lines(buf, 0, -1, false)
+  return find_docstring_end_line_number_in_lines(lines)
+end
+
+M.find_docstring_end_line_number = find_docstring_end_line_number
+
+---Adds an import to the current buffer
 ---@param import_path string: The import path to be added
 ---@param symbol string: The symbol to be imported
-function M.add_import_to_current_buf(import_path, symbol)
-  local docstring_height = find_docstring_end_line_number(0)
-  local insert_on_line
-  if docstring_height == 0 then
-    insert_on_line = 0
-  else
+---@param buf number: The buffer number
+function M.add_import_to_buffer(import_path, symbol, buf)
+  local docstring_height = find_docstring_end_line_number(buf)
+  local insert_on_line = 0
+  if docstring_height ~= 0 then
+    -- add 2 to the docstring height to account for the empty line after the docstring
     insert_on_line = docstring_height + 1
   end
   vim.api.nvim_buf_set_lines(
-    0,
+    buf or 0,
     insert_on_line,
     insert_on_line,
     false,
-    { "from " .. import_path .. " import " .. symbol }
+    { "from " .. import_path .. " import " .. symbol, "" }
   )
 end
 
+---@param path string: The path in which to search for a virtual environment
+---@return string | nil: The path to the virtual environment, or nil if it doesn't exist
+local function dir_contains_virtualenv(path)
+  for _, venv_name in ipairs(cfg.config.python.virtual_env_names) do
+    local venv_path = path .. "/" .. venv_name
+    if vim.fn.isdirectory(venv_path) == 1 then
+      return venv_path
+    end
+  end
+  return nil
+end
+
 ---Get the path to the current virtual environment, or nil if we can't find one
+---@param from_path string: The path to start searching from
 ---@return string | nil: The path to the current virtual environment
-function M.get_virtual_environment()
+function M.get_virtual_environment(from_path)
   local venv = os.getenv("VIRTUAL_ENV")
   if venv then
     return venv
   end
-  local venv_path = vim.fn.getcwd() .. "/.venv"
-  if vim.fn.isdirectory(venv_path) == 1 then
-    return venv_path
+  local current_path = from_path
+  while current_path ~= vim.fn.expand("~") do
+    local venv_path = dir_contains_virtualenv(current_path)
+    if venv_path then
+      return venv_path
+    end
+    current_path = vim.fn.fnamemodify(current_path, ":h")
   end
   return nil
 end
@@ -166,14 +202,14 @@ M.print_msg = print_msg
 ---Print an error message to the console
 ---@param err_msg string: The error message to print
 function M.print_err(err_msg)
-  print_msg(err_msg, config.HL_GROUPS.Error)
+  print_msg(err_msg, cfg.HL_GROUPS.Error)
   log.error(err_msg)
 end
 
 ---Print an info message to the console
 ---@param info_msg string: The info message to print
 function M.print_info(info_msg)
-  print_msg(info_msg, config.HL_GROUPS.More)
+  print_msg(info_msg, cfg.HL_GROUPS.More)
   log.info(info_msg)
 end
 
diff --git a/lua/tests/fixtures/utils/contains_python_files/foo/bar/baz.py b/lua/tests/fixtures/utils/contains_python_files/foo/bar/baz.py
new file mode 100644
index 0000000..e69de29
diff --git a/lua/tests/fixtures/utils/docstrings/misc_1.py b/lua/tests/fixtures/utils/docstrings/misc_1.py
new file mode 100644
index 0000000..bbd84fc
--- /dev/null
+++ b/lua/tests/fixtures/utils/docstrings/misc_1.py
@@ -0,0 +1,4 @@
+"""this is a weirdly formatted docstring
+"""
+
+1 + 1 == 2
diff --git a/lua/tests/fixtures/utils/docstrings/misc_2.py b/lua/tests/fixtures/utils/docstrings/misc_2.py
new file mode 100644
index 0000000..19b37fa
--- /dev/null
+++ b/lua/tests/fixtures/utils/docstrings/misc_2.py
@@ -0,0 +1,4 @@
+"""
+this is a weirdly formatted docstring"""
+
+1 + 1 == 2
diff --git a/lua/tests/fixtures/utils/docstrings/misc_3.py b/lua/tests/fixtures/utils/docstrings/misc_3.py
new file mode 100644
index 0000000..cb01d99
--- /dev/null
+++ b/lua/tests/fixtures/utils/docstrings/misc_3.py
@@ -0,0 +1,6 @@
+"""this 
+is a 
+
+weirdly formatted docstring"""
+
+1 + 1 == 2
diff --git a/lua/tests/fixtures/utils/docstrings/multiline_docstring.py b/lua/tests/fixtures/utils/docstrings/multiline_docstring.py
new file mode 100644
index 0000000..994dbd2
--- /dev/null
+++ b/lua/tests/fixtures/utils/docstrings/multiline_docstring.py
@@ -0,0 +1,9 @@
+"""
+This is a
+multi line docstring
+"""
+
+1 + 1 == 2
+
+if __name__ == "__main__":
+    pass
diff --git a/lua/tests/fixtures/utils/docstrings/no_docstring.py b/lua/tests/fixtures/utils/docstrings/no_docstring.py
new file mode 100644
index 0000000..6c85a60
--- /dev/null
+++ b/lua/tests/fixtures/utils/docstrings/no_docstring.py
@@ -0,0 +1,4 @@
+1 + 1 == 2
+
+if __name__ == "__main__":
+    pass
diff --git a/lua/tests/fixtures/utils/docstrings/single_line_docstring.py b/lua/tests/fixtures/utils/docstrings/single_line_docstring.py
new file mode 100644
index 0000000..2b63801
--- /dev/null
+++ b/lua/tests/fixtures/utils/docstrings/single_line_docstring.py
@@ -0,0 +1,6 @@
+"""This is a single line docstring"""
+
+1 + 1 == 2
+
+if __name__ == "__main__":
+    pass
diff --git a/lua/tests/fixtures/utils/no_python_files/bar/baz.txt b/lua/tests/fixtures/utils/no_python_files/bar/baz.txt
new file mode 100644
index 0000000..e69de29
diff --git a/tests/pymple_spec.lua b/lua/tests/pymple_spec.lua
similarity index 100%
rename from tests/pymple_spec.lua
rename to lua/tests/pymple_spec.lua
diff --git a/lua/tests/resolve_imports_spec.lua b/lua/tests/resolve_imports_spec.lua
new file mode 100644
index 0000000..b658239
--- /dev/null
+++ b/lua/tests/resolve_imports_spec.lua
@@ -0,0 +1,16 @@
+local resolve_imports = require("pymple.resolve_imports")
+
+describe("add_symbol_regexes", function()
+  it("std", function()
+    local args = {}
+    local symbol = "foo"
+    local regexes = { "a %sb", "b %sa" }
+    local result = resolve_imports.add_symbol_regexes(args, symbol, regexes)
+    assert.are.same({
+      "-e",
+      "a foob",
+      "-e",
+      "b fooa",
+    }, result)
+  end)
+end)
diff --git a/lua/tests/update_imports_spec.lua b/lua/tests/update_imports_spec.lua
new file mode 100644
index 0000000..08338f4
--- /dev/null
+++ b/lua/tests/update_imports_spec.lua
@@ -0,0 +1,62 @@
+local update_imports = require("pymple.update_imports")
+local mock = require("luassert.mock")
+
+describe("build_filetypes_args", function()
+  it("std", function()
+    local filetypes = { "python", "lua" }
+    local result = update_imports.build_filetypes_args(filetypes)
+    assert.are.same({ "-t", "python", "-t", "lua" }, result)
+  end)
+end)
+
+describe("update_imports_split", function()
+  it("std", function()
+    local source = "foo/bar/baz"
+    local destination = "oof/rab/zab"
+    local filetypes = { "python", "lua" }
+    local mocked_job = mock(function() end, true)
+    update_imports.update_imports_split(
+      source,
+      destination,
+      filetypes,
+      mocked_job
+    )
+    assert.spy(mocked_job).was_called_with({
+      "--json",
+      "-U",
+      "-t python -t lua",
+      "'from\\s+foo\\.bar\\s+import\\s+\\(?\\n?[\\sa-zA-Z0-9_,\\n]+\\)?\\s*$'",
+      ".",
+    }, { "'%s,%ss/foo\\.bar/oof\\.rab/'" }, true)
+    assert.spy(mocked_job).was_called_with({
+      "--json",
+      "-U",
+      "-t python -t lua",
+      "'from\\s+foo\\.bar\\s+import\\s+\\(?\\n?[\\sa-zA-Z0-9_,\\n]+\\)?\\s*$'",
+      ".",
+    }, { "'%s,%ss/baz/zab/'" }, true)
+    mock.revert(mocked_job)
+  end)
+end)
+
+describe("update_imports_monolithic", function()
+  it("std", function()
+    local source = "foo/bar/baz"
+    local destination = "oof/rab/zab"
+    local filetypes = { "python", "lua" }
+    local mocked_job = mock(function() end, true)
+    update_imports.update_imports_monolithic(
+      source,
+      destination,
+      filetypes,
+      mocked_job
+    )
+    assert.spy(mocked_job).was_called_with({
+      "--json",
+      "-t python -t lua",
+      "'foo\\.bar\\.baz[\\.\\s]'",
+      ".",
+    }, { "'s/foo\\.bar\\.baz\\([\\. ]\\)/oof\\.rab\\.zab\\1/'" }, false)
+    mock.revert(mocked_job)
+  end)
+end)
diff --git a/lua/tests/utils_spec.lua b/lua/tests/utils_spec.lua
new file mode 100644
index 0000000..ee615de
--- /dev/null
+++ b/lua/tests/utils_spec.lua
@@ -0,0 +1,224 @@
+local utils = require("pymple.utils")
+local FIXTURES_PATH = "lua/tests/fixtures/utils"
+local cwd = vim.fn.getcwd()
+local mock = require("luassert.mock")
+
+describe("to_import_path", function()
+  it("std", function()
+    local result = utils.to_import_path("foo/bar/baz.py")
+    assert.equals("foo.bar.baz", result)
+  end)
+end)
+
+describe("split_import_on_last_separator", function()
+  it("std", function()
+    local base_path, module_name =
+      utils.split_import_on_last_separator("foo.bar.baz")
+    assert.equals("foo.bar", base_path)
+    assert.equals("baz", module_name)
+  end)
+
+  it("no separator", function()
+    local base_path, module_name = utils.split_import_on_last_separator("foo")
+    assert.equals("", base_path)
+    assert.equals("foo", module_name)
+  end)
+end)
+
+describe("escape_import_path", function()
+  it("std", function()
+    local result = utils.escape_import_path("foo.bar.baz")
+    assert.equals("foo\\.bar\\.baz", result)
+  end)
+end)
+
+describe("is_python_file", function()
+  it("std", function()
+    local result = utils.is_python_file("foo/bar/baz.py")
+    assert.is_true(result)
+  end)
+
+  it("not python", function()
+    local result = utils.is_python_file("foo/bar/baz.lua")
+    assert.is_false(result)
+  end)
+
+  it("no extension", function()
+    local result = utils.is_python_file("foo/bar/baz")
+    assert.is_false(result)
+  end)
+end)
+
+describe("recursive_dir_contains_python_files", function()
+  it("std", function()
+    local result = utils.recursive_dir_contains_python_files(
+      FIXTURES_PATH .. "/contains_python_files"
+    )
+    assert.is_true(result)
+  end)
+
+  it("no python files", function()
+    local result = utils.recursive_dir_contains_python_files(
+      FIXTURES_PATH .. "/no_python_files"
+    )
+    assert.is_false(result)
+  end)
+end)
+
+describe("find_docstring_end_line_number_in_lines", function()
+  it("no docstring", function()
+    local lines = vim.fn.readfile(
+      cwd .. "/" .. FIXTURES_PATH .. "/docstrings/no_docstring.py"
+    )
+    local result = utils.find_docstring_end_line_number_in_lines(lines)
+    assert.equals(0, result)
+  end)
+
+  it("single line docstring", function()
+    local lines = vim.fn.readfile(
+      cwd .. "/" .. FIXTURES_PATH .. "/docstrings/single_line_docstring.py"
+    )
+    local result = utils.find_docstring_end_line_number_in_lines(lines)
+    assert.equals(1, result)
+  end)
+
+  it("multiline docstring", function()
+    local lines = vim.fn.readfile(
+      cwd .. "/" .. FIXTURES_PATH .. "/docstrings/multiline_docstring.py"
+    )
+    local result = utils.find_docstring_end_line_number_in_lines(lines)
+    assert.equals(4, result)
+  end)
+
+  it("misc 1", function()
+    local lines =
+      vim.fn.readfile(cwd .. "/" .. FIXTURES_PATH .. "/docstrings/misc_1.py")
+    local result = utils.find_docstring_end_line_number_in_lines(lines)
+    assert.equals(2, result)
+  end)
+
+  it("misc 2", function()
+    local lines =
+      vim.fn.readfile(cwd .. "/" .. FIXTURES_PATH .. "/docstrings/misc_2.py")
+    local result = utils.find_docstring_end_line_number_in_lines(lines)
+    assert.equals(2, result)
+  end)
+
+  it("misc 3", function()
+    local lines =
+      vim.fn.readfile(cwd .. "/" .. FIXTURES_PATH .. "/docstrings/misc_3.py")
+    local result = utils.find_docstring_end_line_number_in_lines(lines)
+    assert.equals(4, result)
+  end)
+end)
+
+describe("find_docstring_end_line_number", function()
+  it("std", function()
+    local buf = vim.api.nvim_create_buf(false, true)
+    vim.api.nvim_buf_set_lines(
+      buf,
+      0,
+      -1,
+      false,
+      vim.fn.readfile(
+        cwd .. "/" .. FIXTURES_PATH .. "/docstrings/multiline_docstring.py"
+      )
+    )
+    local result = utils.find_docstring_end_line_number(buf)
+    assert.equals(4, result)
+  end)
+end)
+
+describe("add_import_to_current_buf", function()
+  it("no docstring", function()
+    local buf = vim.api.nvim_create_buf(false, true)
+    vim.api.nvim_buf_set_lines(
+      buf,
+      0,
+      -1,
+      false,
+      vim.fn.readfile(
+        cwd .. "/" .. FIXTURES_PATH .. "/docstrings/no_docstring.py"
+      )
+    )
+    utils.add_import_to_buffer("foo.bar", "baz", buf)
+    local lines = vim.api.nvim_buf_get_lines(buf, 0, -1, false)
+    -- lua table indexing is 1-based
+    assert.equals("from foo.bar import baz", lines[1])
+  end)
+
+  it("single line docstring", function()
+    local buf = vim.api.nvim_create_buf(false, true)
+    vim.api.nvim_buf_set_lines(
+      buf,
+      0,
+      -1,
+      false,
+      vim.fn.readfile(
+        cwd .. "/" .. FIXTURES_PATH .. "/docstrings/single_line_docstring.py"
+      )
+    )
+    utils.add_import_to_buffer("foo.bar", "baz", buf)
+    local lines = vim.api.nvim_buf_get_lines(buf, 0, -1, false)
+    -- lua table indexing is 1-based
+    assert.equals("from foo.bar import baz", lines[3])
+  end)
+
+  it("multiline docstring", function()
+    local buf = vim.api.nvim_create_buf(false, true)
+    vim.api.nvim_buf_set_lines(
+      buf,
+      0,
+      -1,
+      false,
+      vim.fn.readfile(
+        cwd .. "/" .. FIXTURES_PATH .. "/docstrings/multiline_docstring.py"
+      )
+    )
+    utils.add_import_to_buffer("foo.bar", "baz", buf)
+    local lines = vim.api.nvim_buf_get_lines(buf, 0, -1, false)
+    -- lua table indexing is 1-based
+    assert.equals("from foo.bar import baz", lines[6])
+  end)
+end)
+
+describe("get_virtual_environment", function()
+  it("VIRTUAL_ENV set", function()
+    mock(os, "getenv", function()
+      return "/foo/bar"
+    end)
+    local result = utils.get_virtual_environment()
+    assert.equals("/foo/bar", result)
+    mock.revert(os)
+  end)
+
+  it(".venv exists", function()
+    local working_dir = cwd
+      .. "/"
+      .. FIXTURES_PATH
+      .. "/virtual_environments/present"
+    local result = utils.get_virtual_environment(working_dir)
+    assert.equals(working_dir .. "/.venv", result)
+    mock.revert(vim.fn)
+  end)
+
+  it("no venv", function()
+    local working_dir = cwd
+      .. FIXTURES_PATH
+      .. "/virtual_environments/not_present"
+    local result = utils.get_virtual_environment(working_dir)
+    assert.equals(nil, result)
+  end)
+end)
+
+describe("table_contains", function()
+  it("std", function()
+    local result = utils.table_contains({ "foo", "bar", "baz" }, "bar")
+    assert.is_true(result)
+  end)
+
+  it("not present", function()
+    local result = utils.table_contains({ "foo", "bar", "baz" }, "qux")
+    assert.is_false(result)
+  end)
+end)

From 1d2606f7b7b2f929bee13fcc7a25d78b7cbe479d Mon Sep 17 00:00:00 2001
From: Alexandre Pasmantier <alex.pasmant@gmail.com>
Date: Sun, 4 Aug 2024 00:31:27 +0200
Subject: [PATCH 2/4] update docstring

---
 lua/pymple/config.lua | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/lua/pymple/config.lua b/lua/pymple/config.lua
index 7910a7c..d3fc844 100644
--- a/lua/pymple/config.lua
+++ b/lua/pymple/config.lua
@@ -54,7 +54,12 @@
 ---     -- the log level to use
 ---     -- (one of "trace", "debug", "info", "warn", "error", "fatal")
 ---     level = "debug",
----   }
+---   },
+---   -- python options
+---   python = {
+---   -- the names of virtual environment folders to look out for
+---   virtual_env_names = { ".venv" },
+---   },
 --- }
 --- ```
 --- </pre>

From 66393b90b2db9ebb1ae6dda043365262a443ff04 Mon Sep 17 00:00:00 2001
From: Github Actions <actions@github>
Date: Sat, 3 Aug 2024 22:31:42 +0000
Subject: [PATCH 3/4] [docgen] Update doc/pymple.txt skip-checks: true

---
 doc/pymple.txt | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/doc/pymple.txt b/doc/pymple.txt
index a221044..d546e7b 100644
--- a/doc/pymple.txt
+++ b/doc/pymple.txt
@@ -77,7 +77,12 @@ config = {
     -- the log level to use
     -- (one of "trace", "debug", "info", "warn", "error", "fatal")
     level = "debug",
-  }
+  },
+  -- python options
+  python = {
+  -- the names of virtual environment folders to look out for
+  virtual_env_names = { ".venv" },
+  },
 }
 ```
 

From d206dc336ed59e791d991f7bcecb1c61c7d2805a Mon Sep 17 00:00:00 2001
From: Alexandre Pasmantier <alex.pasmant@gmail.com>
Date: Sun, 4 Aug 2024 00:36:40 +0200
Subject: [PATCH 4/4] add virtualenvs fixtures

---
 .../fixtures/utils/virtual_environments/not_present/empty.txt     | 0
 .../fixtures/utils/virtual_environments/present/.venv/empty.txt   | 0
 2 files changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 lua/tests/fixtures/utils/virtual_environments/not_present/empty.txt
 create mode 100644 lua/tests/fixtures/utils/virtual_environments/present/.venv/empty.txt

diff --git a/lua/tests/fixtures/utils/virtual_environments/not_present/empty.txt b/lua/tests/fixtures/utils/virtual_environments/not_present/empty.txt
new file mode 100644
index 0000000..e69de29
diff --git a/lua/tests/fixtures/utils/virtual_environments/present/.venv/empty.txt b/lua/tests/fixtures/utils/virtual_environments/present/.venv/empty.txt
new file mode 100644
index 0000000..e69de29