diff --git a/TODO.md b/TODO.md
index 1d11051..3fd1422 100644
--- a/TODO.md
+++ b/TODO.md
@@ -5,3 +5,5 @@
     * Duh just check to see it has no parents
 * Swapping
 * Jumplist
+* Get more languages into movement/highlight/swap specs
+    * :help treesitter-parsers
diff --git a/tests/fixtures/python.py b/tests/fixtures/python.py
index 04b9ae7..0328cf1 100644
--- a/tests/fixtures/python.py
+++ b/tests/fixtures/python.py
@@ -1,5 +1,14 @@
 from functools import wraps
 
+def random_annotation(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        # Perform some task related to randomness
+        print("Random thing happening")
+        return func(*args, **kwargs)
+    return wrapper
+
+
 class Person:
     def __init__(self, name):
         self.name = name
@@ -15,6 +24,7 @@ def __init__(self, title, author):
     def describe(self):
         print(f"{self.title} by {self.author}")
 
+@random_annotation
 class Car:
     def __init__(self, make, model, year):
         self.make = make
@@ -25,15 +35,6 @@ def display_info(self):
         print(f"Make: {self.make}, Model: {self.model}, Year: {self.year}")
 
 
-def random_annotation(func):
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        # Perform some task related to randomness
-        print("Random thing happening")
-        return func(*args, **kwargs)
-    return wrapper
-
-
 @random_annotation
 def main():
     """
diff --git a/tests/load_fixture.lua b/tests/load_fixture.lua
index 57f6d9c..36ec710 100644
--- a/tests/load_fixture.lua
+++ b/tests/load_fixture.lua
@@ -1,11 +1,16 @@
 local fixtures_dir = vim.fn.expand 'tests/fixtures'
 
+---@param filename string
+local function get_file_extension(filename)
+  return filename:match(".+%.([^.]+)$")
+end
+
 -- Creates a buffer, and loads the contents of a specified filename into it.
 -- Sets it as the current buffer
 ---@param filename string
----@param lang string same as filename extension, laziness was here
 ---@return integer
-local function load_fixture(filename, lang)
+local function load_fixture(filename)
+  local lang = get_file_extension(filename)
   local buf = vim.api.nvim_create_buf(false, true) -- Create a new buffer (listed) to enable interaction with it
   local lines = {}
 
diff --git a/tests/treewalker/highlight_spec.lua b/tests/treewalker/highlight_spec.lua
index b8fa25d..31b3fd4 100644
--- a/tests/treewalker/highlight_spec.lua
+++ b/tests/treewalker/highlight_spec.lua
@@ -18,7 +18,7 @@ describe("Treewalker highlighting", function()
   end
 
   describe("regular lua file: ", function()
-    load_fixture("/lua.lua", "lua")
+    load_fixture("/lua.lua")
 
     before_each(function()
       treewalker.setup({ highlight = true })
diff --git a/tests/treewalker/movement_spec.lua b/tests/treewalker/movement_spec.lua
index b8f7a19..ec7e644 100644
--- a/tests/treewalker/movement_spec.lua
+++ b/tests/treewalker/movement_spec.lua
@@ -20,7 +20,7 @@ end
 
 describe("Treewalker movement", function()
   describe("regular lua file: ", function()
-    load_fixture("/lua.lua", "lua")
+    load_fixture("/lua.lua")
 
     it("moves up and down at the same pace", function()
       vim.fn.cursor(1, 1) -- Reset cursor
@@ -80,7 +80,7 @@ describe("Treewalker movement", function()
   end)
 
   describe("lua spec file: ", function()
-    load_fixture("/lua-spec.lua", "lua")
+    load_fixture("/lua-spec.lua")
 
     -- go to first describe
     local function go_to_describe()