From 42bc5353e3de9a8a0c8fdcafd7ef4ee7534105ac Mon Sep 17 00:00:00 2001
From: Max Schaefer <max@semmle.com>
Date: Mon, 6 Apr 2020 14:45:59 +0100
Subject: [PATCH] Refine our modelling of test functions and split it out into
 a separate library.

---
 ql/src/filters/ClassifyFiles.ql         | 11 +++----
 ql/src/go.qll                           |  1 +
 ql/src/semmle/go/Decls.qll              |  7 +++++
 ql/src/semmle/go/frameworks/Testing.qll | 42 +++++++++++++++++++++++++
 4 files changed, 54 insertions(+), 7 deletions(-)
 create mode 100644 ql/src/semmle/go/frameworks/Testing.qll

diff --git a/ql/src/filters/ClassifyFiles.ql b/ql/src/filters/ClassifyFiles.ql
index 469ed63033c..120ee3100a8 100644
--- a/ql/src/filters/ClassifyFiles.ql
+++ b/ql/src/filters/ClassifyFiles.ql
@@ -12,19 +12,16 @@ string generatorCommentRegex() {
   result = "Generated By\\b.*\\bDo not edit" or
   result = "This (file|class|interface|art[ei]fact) (was|is|(has been)) (?:auto[ -]?)?gener(e?)ated" or
   result = "Any modifications to this file will be lost" or
-  result = "This (file|class|interface|art[ei]fact) (was|is) (?:mechanically|automatically) generated" or
+  result =
+    "This (file|class|interface|art[ei]fact) (was|is) (?:mechanically|automatically) generated" or
   result = "The following code was (?:auto[ -]?)?generated (?:by|from)" or
   result = "Autogenerated by Thrift" or
   result = "(Code g|G)enerated from .* by ANTLR"
 }
 
 predicate classify(File f, string category) {
-  // `go test`-style test
-  f.getBaseName().regexpMatch(".*_test.go") and
-  exists(FuncDecl fn |
-    fn.getName().regexpMatch("(Test|Benchmark|Example)[^a-z].*") and
-    fn.getFile() = f
-  ) and
+  // tests
+  f = any(TestCase tc).getFile() and
   category = "test"
   or
   // vendored code
diff --git a/ql/src/go.qll b/ql/src/go.qll
index c176857a282..0bce4c7d6b7 100644
--- a/ql/src/go.qll
+++ b/ql/src/go.qll
@@ -28,5 +28,6 @@ import semmle.go.frameworks.SystemCommandExecutors
 import semmle.go.frameworks.SQL
 import semmle.go.frameworks.XPath
 import semmle.go.frameworks.Stdlib
+import semmle.go.frameworks.Testing
 import semmle.go.security.FlowSources
 import semmle.go.Util
diff --git a/ql/src/semmle/go/Decls.qll b/ql/src/semmle/go/Decls.qll
index c8ce8f6047a..8e8861247d4 100644
--- a/ql/src/semmle/go/Decls.qll
+++ b/ql/src/semmle/go/Decls.qll
@@ -129,6 +129,13 @@ class FuncDef extends @funcdef, StmtParent, ExprParent {
     result.getFunction() = this
   }
 
+  /**
+   * Gets the number of parameters of this function.
+   */
+  int getNumParameter() {
+    result = count(getAParameter())
+  }
+
   /**
    * Gets a call to this function.
    */
diff --git a/ql/src/semmle/go/frameworks/Testing.qll b/ql/src/semmle/go/frameworks/Testing.qll
new file mode 100644
index 00000000000..33548a8720e
--- /dev/null
+++ b/ql/src/semmle/go/frameworks/Testing.qll
@@ -0,0 +1,42 @@
+/** Provides classes for working with tests. */
+
+import go
+
+/**
+ * A program element that represents a test case.
+ *
+ * Extend this class to refine existing models of testing frameworks. If you want to model new
+ * frameworks, extend `TestCase::Range` instead.
+ */
+class TestCase extends AstNode {
+  TestCase::Range self;
+
+  TestCase() { this = self }
+}
+
+/** Provides classes for working with test cases. */
+module TestCase {
+  /**
+   * A program element that represents a test case.
+   *
+   * Extend this class to model new testing frameworks. If you want to refine existing models,
+   * extend `TestCase` instead.
+   */
+  abstract class Range extends AstNode { }
+
+  /** A `go test` style test (including benchmarks and examples). */
+  private class GoTestFunction extends Range, FuncDef {
+    GoTestFunction() {
+      getName().regexpMatch("Test[^a-z].*") and
+      getNumParameter() = 1 and
+      getParameter(0).getType().(PointerType).getBaseType().hasQualifiedName("testing", "T")
+      or
+      getName().regexpMatch("Benchmark[^a-z].*") and
+      getNumParameter() = 1 and
+      getParameter(0).getType().(PointerType).getBaseType().hasQualifiedName("testing", "B")
+      or
+      getName().regexpMatch("Example[^a-z].*") and
+      getNumParameter() = 0
+    }
+  }
+}