From 2dca2802860c802dd0ca2eb1e0057f96828debfb Mon Sep 17 00:00:00 2001
From: Michael Noah <92764374+mnoah1@users.noreply.github.com>
Date: Tue, 2 Jul 2024 13:38:59 -0400
Subject: [PATCH] Adjustments to Python test discovery (#17)

Making a few tweaks to improve Python test case discovery in files,
based on some initial user feedback:
- Use selection range for positioning of arrows, so arrow lines up with
the function header. The previous range value began at the start of any
decorators or other content above the header.
  Example of adjusted positions:

![image](https://github.com/uber/vscode-bazel-bsp/assets/92764374/ec3096f8-4ab3-420b-9f18-20cccc8551f9)

- Let class names start _or_ end with Test
- Add a filter before returning that excludes classes containing no test
cases. This will reduce noise in cases where the word Test is in a class
name but it is not in fact a test.
---
 src/language-tools/python.ts                 | 14 +++--
 src/test/suite/language-tools/python.test.ts | 54 +++++++++++++++++++-
 2 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/src/language-tools/python.ts b/src/language-tools/python.ts
index 1e2e9dd..2832300 100644
--- a/src/language-tools/python.ts
+++ b/src/language-tools/python.ts
@@ -88,13 +88,14 @@ export class PythonLanguageTools
         document
       )
     const result: DocumentTestItem[] = []
+    const shouldInclude = new Set<DocumentTestItem>()
     const evaluateCurrentSymbol = (
       symbol: vscode.DocumentSymbol,
       parent: DocumentTestItem | undefined = undefined
     ) => {
       let newItem: DocumentTestItem | undefined
       if (symbol.kind === vscode.SymbolKind.Class) {
-        if (symbol.name.startsWith('Test')) {
+        if (symbol.name.startsWith('Test') || symbol.name.endsWith('Test')) {
           // Capture class names that begin with 'Test'
           newItem = {
             name: symbol.name,
@@ -124,12 +125,15 @@ export class PythonLanguageTools
 
         newItem = {
           name: symbol.name,
-          range: symbol.range,
+          range: symbol.selectionRange,
           uri: document,
           testFilter: testFilter,
           parent: parent,
           lookupKey: lookupKey,
         }
+
+        if (parent) shouldInclude.add(parent)
+        shouldInclude.add(newItem)
       }
       if (newItem) {
         result.push(newItem)
@@ -146,9 +150,13 @@ export class PythonLanguageTools
       evaluateCurrentSymbol(symbol)
     }
 
+    const finalTestCases = result.filter(item => {
+      return shouldInclude.has(item)
+    })
+
     return {
       isTestFile: true,
-      testCases: result,
+      testCases: finalTestCases,
       documentTest: fullDocTestItem,
     }
   }
diff --git a/src/test/suite/language-tools/python.test.ts b/src/test/suite/language-tools/python.test.ts
index 66d4da4..6c510de 100644
--- a/src/test/suite/language-tools/python.test.ts
+++ b/src/test/suite/language-tools/python.test.ts
@@ -36,7 +36,7 @@ suite('Python Language Tools', () => {
       '/repo/root/'
     )
     assert.strictEqual(result.isTestFile, true)
-    assert.strictEqual(result.testCases.length, 5)
+    assert.strictEqual(result.testCases.length, 7)
 
     const expectedTests: Map<string, string[]> = new Map()
     expectedTests.set('TestExample', [
@@ -59,6 +59,14 @@ suite('Python Language Tools', () => {
       'sample.my_test.test_separate_function',
       'my_test and test_separate_function',
     ])
+    expectedTests.set('ExampleTest', [
+      'sample.my_test.ExampleTest',
+      'my_test and ExampleTest',
+    ])
+    expectedTests.set('test_example_2', [
+      'sample.my_test.ExampleTest.test_example_2',
+      'my_test and ExampleTest and test_example_2',
+    ])
 
     for (const test of result.testCases) {
       const expected = expectedTests.get(test.name)
@@ -257,4 +265,48 @@ const sampleDocumentSymbols: vscode.DocumentSymbol[] = [
     ),
     children: [],
   },
+  {
+    // Class name that matches but contains no methods.
+    name: 'TestExampleEmpty',
+    detail: '',
+    kind: vscode.SymbolKind.Class,
+    range: new vscode.Range(
+      new vscode.Position(8, 0),
+      new vscode.Position(35, 44)
+    ),
+    selectionRange: new vscode.Range(
+      new vscode.Position(8, 6),
+      new vscode.Position(8, 25)
+    ),
+    children: [],
+  },
+  {
+    name: 'ExampleTest',
+    detail: '',
+    kind: vscode.SymbolKind.Class,
+    range: new vscode.Range(
+      new vscode.Position(8, 0),
+      new vscode.Position(35, 44)
+    ),
+    selectionRange: new vscode.Range(
+      new vscode.Position(8, 6),
+      new vscode.Position(8, 25)
+    ),
+    children: [
+      {
+        name: 'test_example_2',
+        detail: '',
+        kind: vscode.SymbolKind.Method,
+        range: new vscode.Range(
+          new vscode.Position(10, 4),
+          new vscode.Position(12, 5)
+        ),
+        selectionRange: new vscode.Range(
+          new vscode.Position(10, 4),
+          new vscode.Position(10, 16)
+        ),
+        children: [],
+      },
+    ],
+  },
 ]