1; subclasses of unittest.TestCase or TestCase
  2(
  3    (class_definition
  4        name: (identifier) @run @_unittest_class_name
  5        superclasses: (argument_list
  6            [(identifier) @_superclass
  7                (attribute (identifier) @_superclass)]
  8            )
  9        (#eq? @_superclass "TestCase")
 10        ) @_python-unittest-class
 11    (#set! tag python-unittest-class)
 12    )
 13
 14; test methods whose names start with `test` in a TestCase
 15(
 16    (class_definition
 17        name: (identifier) @_unittest_class_name
 18        superclasses: (argument_list
 19            [(identifier) @_superclass
 20                (attribute (identifier) @_superclass)]
 21            )
 22        (#eq? @_superclass "TestCase")
 23        body: (block
 24            (function_definition
 25                name: (identifier) @run @_unittest_method_name
 26                (#match? @_unittest_method_name "^test.*")
 27                ) @_python-unittest-method
 28            (#set! tag python-unittest-method)
 29            )
 30        )
 31    )
 32
 33; pytest functions
 34(
 35    (module
 36        (function_definition
 37            name: (identifier) @run @_pytest_method_name
 38            (#match? @_pytest_method_name "^test_")
 39            ) @_python-pytest-method
 40        )
 41    (#set! tag python-pytest-method)
 42    )
 43
 44; decorated pytest functions
 45(
 46    (module
 47        (decorated_definition
 48            (decorator)+ @_decorator
 49            definition: (function_definition
 50                name: (identifier) @run @_pytest_method_name
 51                (#match? @_pytest_method_name "^test_")
 52                )
 53            ) @_python-pytest-method
 54        )
 55    (#set! tag python-pytest-method)
 56    )
 57
 58
 59; pytest classes
 60(
 61    (module
 62        (class_definition
 63            name: (identifier) @run @_pytest_class_name
 64            (#match? @_pytest_class_name "^Test")
 65            )
 66        (#set! tag python-pytest-class)
 67        )
 68    )
 69
 70
 71; decorated pytest classes
 72(
 73    (module
 74        (decorated_definition
 75            (decorator)+ @_decorator
 76            definition: (class_definition
 77                name: (identifier) @run @_pytest_class_name
 78                (#match? @_pytest_class_name "^Test")
 79                )
 80            )
 81        (#set! tag python-pytest-class)
 82        )
 83    )
 84
 85
 86; pytest class methods
 87(
 88    (module
 89        (class_definition
 90            name: (identifier) @_pytest_class_name
 91            (#match? @_pytest_class_name "^Test")
 92            body: (block
 93                [(decorated_definition
 94                    (decorator)+ @_decorator
 95                    definition: (function_definition
 96                        name: (identifier) @run @_pytest_method_name
 97                        (#match? @_pytest_method_name "^test_")
 98                        )
 99                    )
100                (function_definition
101                    name: (identifier) @run @_pytest_method_name
102                    (#match? @_pytest_method_name "^test")
103                    )
104                ] @_python-pytest-method)
105            (#set! tag python-pytest-method)
106            )
107        )
108    )
109
110; decorated pytest class methods
111(
112    (module
113        (decorated_definition
114            (decorator)+ @_decorator
115            definition: (class_definition
116                name: (identifier) @_pytest_class_name
117                (#match? @_pytest_class_name "^Test")
118                body: (block
119                    [(decorated_definition
120                        (decorator)+ @_decorator
121                        definition: (function_definition
122                            name: (identifier) @run @_pytest_method_name
123                            (#match? @_pytest_method_name "^test_")
124                            )
125                        )
126                    (function_definition
127                        name: (identifier) @run @_pytest_method_name
128                        (#match? @_pytest_method_name "^test")
129                        )
130                    ] @_python-pytest-method)
131                (#set! tag python-pytest-method)
132                )
133            )
134        )
135    )
136
137; module main method
138(
139    (module
140        (if_statement
141            condition: (comparison_operator
142                (identifier) @run @_lhs
143                operators: "=="
144                (string) @_rhs
145                )
146            (#eq? @_lhs "__name__")
147            (#match? @_rhs "^[\"']__main__[\"']$")
148            (#set! tag python-module-main-method)
149            )
150        )
151    )