From 16c3c90b552f56b8817c4c42bee9779d4f683c24 Mon Sep 17 00:00:00 2001
From: Daniel Watkins <daniel@daniel-watkins.co.uk>
Date: Thu, 10 Nov 2022 03:41:54 -0500
Subject: [PATCH] feat(query): include test decorators in the test definition
 (#34)

This means that `neotest.run.run` when the cursor is on a test's
decorators will run the decorated test (instead of running the previous
test in the file, the current behaviour).

Fixes: #31
---
 lua/neotest-python/init.lua | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/lua/neotest-python/init.lua b/lua/neotest-python/init.lua
index 900e91b..0847187 100644
--- a/lua/neotest-python/init.lua
+++ b/lua/neotest-python/init.lua
@@ -83,14 +83,31 @@ end
 ---@return Tree | nil
 function PythonNeotestAdapter.discover_positions(path)
   local query = [[
+    ;; Match undecorated functions
     ((function_definition
       name: (identifier) @test.name)
       (#match? @test.name "^test"))
       @test.definition
+    ;; Match decorated function, including decorators in definition
+    (decorated_definition
+      ((function_definition
+        name: (identifier) @test.name)
+        (#match? @test.name "^test")))
+        @test.definition
 
-    (class_definition
-     name: (identifier) @namespace.name)
-     @namespace.definition
+    ;; Match decorated classes, including decorators in definition
+    (decorated_definition
+      (class_definition
+       name: (identifier) @namespace.name))
+      @namespace.definition
+    ;; Match undecorated classes: namespaces nest so #not-has-parent is used
+    ;; to ensure each namespace is annotated only once
+    (
+     (class_definition
+      name: (identifier) @namespace.name)
+      @namespace.definition
+     (#not-has-parent? @namespace.definition decorated_definition)
+    )
   ]]
   local root = PythonNeotestAdapter.root(path)
   local python = get_python(root)