; subclasses of unittest.TestCase or TestCase
((class_definition
  name: (identifier) @run @_unittest_class_name
  superclasses: (argument_list
    [
      (identifier) @_superclass
      (attribute
        (identifier) @_superclass)
    ])
  (#eq? @_superclass "TestCase")) @_python-unittest-class
  (#set! tag python-unittest-class))

; test methods whose names start with `test` in a TestCase
(class_definition
  name: (identifier) @_unittest_class_name
  superclasses: (argument_list
    [
      (identifier) @_superclass
      (attribute
        (identifier) @_superclass)
    ])
  (#eq? @_superclass "TestCase")
  body: (block
    (function_definition
      name: (identifier) @run @_unittest_method_name
      (#match? @_unittest_method_name "^test.*")) @_python-unittest-method
    (#set! tag python-unittest-method)))

; pytest functions
((module
  (function_definition
    name: (identifier) @run @_pytest_method_name
    (#match? @_pytest_method_name "^test_")) @_python-pytest-method)
  (#set! tag python-pytest-method))

; decorated pytest functions
((module
  (decorated_definition
    (decorator)+ @_decorator
    definition: (function_definition
      name: (identifier) @run @_pytest_method_name
      (#match? @_pytest_method_name "^test_"))) @_python-pytest-method)
  (#set! tag python-pytest-method))

; pytest classes
(module
  (class_definition
    name: (identifier) @run @_pytest_class_name
    (#match? @_pytest_class_name "^Test"))
  (#set! tag python-pytest-class))

; decorated pytest classes
(module
  (decorated_definition
    (decorator)+ @_decorator
    definition: (class_definition
      name: (identifier) @run @_pytest_class_name
      (#match? @_pytest_class_name "^Test")))
  (#set! tag python-pytest-class))

; pytest class methods
(module
  (class_definition
    name: (identifier) @_pytest_class_name
    (#match? @_pytest_class_name "^Test")
    body: (block
      [
        (decorated_definition
          (decorator)+ @_decorator
          definition: (function_definition
            name: (identifier) @run @_pytest_method_name
            (#match? @_pytest_method_name "^test_")))
        (function_definition
          name: (identifier) @run @_pytest_method_name
          (#match? @_pytest_method_name "^test"))
      ] @_python-pytest-method)
    (#set! tag python-pytest-method)))

; decorated pytest class methods
(module
  (decorated_definition
    (decorator)+ @_decorator
    definition: (class_definition
      name: (identifier) @_pytest_class_name
      (#match? @_pytest_class_name "^Test")
      body: (block
        [
          (decorated_definition
            (decorator)+ @_decorator
            definition: (function_definition
              name: (identifier) @run @_pytest_method_name
              (#match? @_pytest_method_name "^test_")))
          (function_definition
            name: (identifier) @run @_pytest_method_name
            (#match? @_pytest_method_name "^test"))
        ] @_python-pytest-method)
      (#set! tag python-pytest-method))))

; module main method
(module
  (if_statement
    condition: (comparison_operator
      (identifier) @run @_lhs
      operators: "=="
      (string) @_rhs)
    (#eq? @_lhs "__name__")
    (#match? @_rhs "^[\"']__main__[\"']$")
    (#set! tag python-module-main-method)))