import python

private import semmle.python.ApiGraphs

abstract class TestScope extends Scope { }

class UnitTestClass extends TestScope, Class {

UnitTestClass() {

exists(API::Node testCaseClass, string testCaseString |

testCaseString.matches("%TestCase") and

testCaseClass = any(API::Node mod).getMember(testCaseString)

|

this.getParent() = testCaseClass.getASubclass*().asSource().asExpr()

)

}

}

abstract class Test extends TestScope { }

/** A test function that uses the `unittest` framework */

class UnitTestFunction extends Test {

UnitTestFunction() {

this.getScope+() instanceof UnitTestClass and

this.(Function).getName().matches("test%")

}

}

class PyTestFunction extends Test {

PyTestFunction() {

exists(Module pytest | pytest.getName() = "pytest") and

this.(Function).getName().matches("test%")

}

}

class NoseTestFunction extends Test {

NoseTestFunction() {

exists(Module nose | nose.getName() = "nose") and

this.(Function).getName().matches("test%")

}

}

/** A function that is clearly a test, but doesn't belong to a specific framework */

class UnknownTestFunction extends Test {

UnknownTestFunction() {

this.(Function).getName().matches("test%") and

this.getEnclosingModule().getFile().getShortName().matches("test_%.py")

}

}