Add location to TestClass nodes for Test Explorer runnability (pytest and unittest)

python_files/tests/pytestadapter/expected_discovery_test_output.py

@@ -1,6 +1,11 @@
 import os
 
-from .helpers import TEST_DATA_PATH, find_test_line_number, get_absolute_test_id
+from .helpers import (
+    TEST_DATA_PATH,
+    find_class_line_number,
+    find_test_line_number,
+    get_absolute_test_id,
+)
 
 # This file contains the expected output dictionaries for tests discovery and is used in test_discovery.py.
 
@@ -95,6 +100,7 @@
                         "unittest_pytest_same_file.py::TestExample",
                         unit_pytest_same_file_path,
                     ),
+                    "lineno": find_class_line_number("TestExample", unit_pytest_same_file_path),
                 },
                 {
                     "name": "test_true_pytest",
@@ -207,6 +213,7 @@
                                 "unittest_folder/test_add.py::TestAddFunction",
                                 test_add_path,
                             ),
+                            "lineno": find_class_line_number("TestAddFunction", test_add_path),
                         },
                         {
                             "name": "TestDuplicateFunction",
@@ -235,6 +242,9 @@
                                 "unittest_folder/test_add.py::TestDuplicateFunction",
                                 test_add_path,
                             ),
+                            "lineno": find_class_line_number(
+                                "TestDuplicateFunction", test_add_path
+                            ),
                         },
                     ],
                 },
@@ -288,6 +298,9 @@
                                 "unittest_folder/test_subtract.py::TestSubtractFunction",
                                 test_subtract_path,
                             ),
+                            "lineno": find_class_line_number(
+                                "TestSubtractFunction", test_subtract_path
+                            ),
                         },
                         {
                             "name": "TestDuplicateFunction",
@@ -316,6 +329,9 @@
                                 "unittest_folder/test_subtract.py::TestDuplicateFunction",
                                 test_subtract_path,
                             ),
+                            "lineno": find_class_line_number(
+                                "TestDuplicateFunction", test_subtract_path
+                            ),
                         },
                     ],
                 },
@@ -553,6 +569,7 @@
                         "parametrize_tests.py::TestClass",
                         parameterize_tests_path,
                     ),
+                    "lineno": find_class_line_number("TestClass", parameterize_tests_path),
                     "children": [
                         {
                             "name": "test_adding",
@@ -929,6 +946,7 @@
                         "test_multi_class_nest.py::TestFirstClass",
                         TEST_MULTI_CLASS_NEST_PATH,
                     ),
+                    "lineno": find_class_line_number("TestFirstClass", TEST_MULTI_CLASS_NEST_PATH),
                     "children": [
                         {
                             "name": "TestSecondClass",
@@ -938,6 +956,9 @@
                                 "test_multi_class_nest.py::TestFirstClass::TestSecondClass",
                                 TEST_MULTI_CLASS_NEST_PATH,
                             ),
+                            "lineno": find_class_line_number(
+                                "TestSecondClass", TEST_MULTI_CLASS_NEST_PATH
+                            ),
                             "children": [
                                 {
                                     "name": "test_second",
@@ -982,6 +1003,9 @@
                                 "test_multi_class_nest.py::TestFirstClass::TestSecondClass2",
                                 TEST_MULTI_CLASS_NEST_PATH,
                             ),
+                            "lineno": find_class_line_number(
+                                "TestSecondClass2", TEST_MULTI_CLASS_NEST_PATH
+                            ),
                             "children": [
                                 {
                                     "name": "test_second2",
@@ -1227,6 +1251,9 @@
                         "same_function_new_class_param.py::TestNotEmpty",
                         TEST_DATA_PATH / "same_function_new_class_param.py",
                     ),
+                    "lineno": find_class_line_number(
+                        "TestNotEmpty", TEST_DATA_PATH / "same_function_new_class_param.py"
+                    ),
                 },
                 {
                     "name": "TestEmpty",
@@ -1298,6 +1325,9 @@
                         "same_function_new_class_param.py::TestEmpty",
                         TEST_DATA_PATH / "same_function_new_class_param.py",
                     ),
+                    "lineno": find_class_line_number(
+                        "TestEmpty", TEST_DATA_PATH / "same_function_new_class_param.py"
+                    ),
                 },
             ],
         }
@@ -1371,6 +1401,9 @@
                         "test_param_span_class.py::TestClass1",
                         TEST_DATA_PATH / "test_param_span_class.py",
                     ),
+                    "lineno": find_class_line_number(
+                        "TestClass1", TEST_DATA_PATH / "test_param_span_class.py"
+                    ),
                 },
                 {
                     "name": "TestClass2",
@@ -1427,6 +1460,9 @@
                         "test_param_span_class.py::TestClass2",
                         TEST_DATA_PATH / "test_param_span_class.py",
                     ),
+                    "lineno": find_class_line_number(
+                        "TestClass2", TEST_DATA_PATH / "test_param_span_class.py"
+                    ),
                 },
             ],
         }
@@ -1503,6 +1539,7 @@
                                 "pytest_describe_plugin/describe_only.py::describe_A",
                                 describe_only_path,
                             ),
+                            "lineno": find_class_line_number("describe_A", describe_only_path),
                         }
                     ],
                 }
@@ -1586,6 +1623,9 @@
                                         "pytest_describe_plugin/nested_describe.py::describe_list::describe_append",
                                         nested_describe_path,
                                     ),
+                                    "lineno": find_class_line_number(
+                                        "describe_append", nested_describe_path
+                                    ),
                                 },
                                 {
                                     "name": "describe_remove",
@@ -1614,12 +1654,16 @@
                                         "pytest_describe_plugin/nested_describe.py::describe_list::describe_remove",
                                         nested_describe_path,
                                     ),
+                                    "lineno": find_class_line_number(
+                                        "describe_remove", nested_describe_path
+                                    ),
                                 },
                             ],
                             "id_": get_absolute_test_id(
                                 "pytest_describe_plugin/nested_describe.py::describe_list",
                                 nested_describe_path,
                             ),
+                            "lineno": find_class_line_number("describe_list", nested_describe_path),
                         }
                     ],
                 }

python_files/tests/pytestadapter/helpers.py

@@ -370,6 +370,28 @@ def find_test_line_number(test_name: str, test_file_path) -> str:
     raise ValueError(error_str)
 
 
+def find_class_line_number(class_name: str, test_file_path) -> str:
+    """Function which finds the correct line number for a class definition.
+
+    Args:
+    class_name: The name of the class to find the line number for.
+    test_file_path: The path to the test file where the class is located.
+    """
+    # Look for the class definition line (or function for pytest-describe)
+    with open(test_file_path) as f:  # noqa: PTH123
+        for i, line in enumerate(f):
+            # Match "class ClassName" or "class ClassName(" or "class ClassName:"
+            # Also match "def ClassName(" for pytest-describe blocks
+            if (
+                line.strip().startswith(f"class {class_name}")
+                or line.strip().startswith(f"class {class_name}(")
+                or line.strip().startswith(f"def {class_name}(")
+            ):
+                return str(i + 1)
+    error_str: str = f"Class {class_name!r} not found on any line in {test_file_path}"
+    raise ValueError(error_str)
+
+
 def get_absolute_test_id(test_id: str, test_path: pathlib.Path) -> str:
     """Get the absolute test id by joining the testPath with the test_id."""
     split_id = test_id.split("::")[1:]

python_files/tests/unittestadapter/expected_discovery_test_output.py

@@ -9,6 +9,25 @@
 TEST_DATA_PATH = pathlib.Path(__file__).parent / ".data"
 
 
+def find_class_line_number(class_name: str, test_file_path) -> str:
+    """Function which finds the correct line number for a class definition.
+
+    Args:
+    class_name: The name of the class to find the line number for.
+    test_file_path: The path to the test file where the class is located.
+    """
+    # Look for the class definition line
+    with pathlib.Path(test_file_path).open() as f:
+        for i, line in enumerate(f):
+            # Match "class ClassName" or "class ClassName(" or "class ClassName:"
+            if line.strip().startswith(f"class {class_name}") or line.strip().startswith(
+                f"class {class_name}("
+            ):
+                return str(i + 1)
+    error_str: str = f"Class {class_name!r} not found on any line in {test_file_path}"
+    raise ValueError(error_str)
+
+
 skip_unittest_folder_discovery_output = {
     "path": os.fspath(TEST_DATA_PATH / "unittest_skip"),
     "name": "unittest_skip",
@@ -49,6 +68,10 @@
                     ],
                     "id_": os.fspath(TEST_DATA_PATH / "unittest_skip" / "unittest_skip_function.py")
                     + "\\SimpleTest",
+                    "lineno": find_class_line_number(
+                        "SimpleTest",
+                        TEST_DATA_PATH / "unittest_skip" / "unittest_skip_function.py",
+                    ),
                 }
             ],
             "id_": os.fspath(TEST_DATA_PATH / "unittest_skip" / "unittest_skip_function.py"),
@@ -114,6 +137,16 @@
                                         },
                                     ],
                                     "id_": complex_tree_file_path + "\\" + "TreeOne",
+                                    "lineno": find_class_line_number(
+                                        "TreeOne",
+                                        pathlib.PurePath(
+                                            TEST_DATA_PATH,
+                                            "utils_complex_tree",
+                                            "test_outer_folder",
+                                            "test_inner_folder",
+                                            "test_utils_complex_tree.py",
+                                        ),
+                                    ),
                                 }
                             ],
                             "id_": complex_tree_file_path,

python_files/unittestadapter/pvsc_utils.py

@@ -44,6 +44,7 @@ class TestItem(TestData):
 
 class TestNode(TestData):
     children: "List[TestNode | TestItem]"
+    lineno: NotRequired[str]  # Optional field for class nodes
 
 
 class TestExecutionStatus(str, enum.Enum):
@@ -101,6 +102,16 @@ def get_test_case(suite):
             yield from get_test_case(test)
 
 
+def get_class_line(test_case: unittest.TestCase) -> str:
+    """Get the line number where a test class is defined."""
+    try:
+        test_class = test_case.__class__
+        _sourcelines, lineno = inspect.getsourcelines(test_class)
+        return str(lineno)
+    except Exception:
+        return ""
+
+
 def get_source_line(obj) -> str:
     """Get the line number of a test case start line."""
     try:
@@ -249,6 +260,10 @@ def build_test_tree(
                 class_name, file_path, TestNodeTypeEnum.class_, current_node
             )
 
+            # Add line number to class node if not already present.
+            if "lineno" not in current_node:
+                current_node["lineno"] = get_class_line(test_case)
+
             # Get test line number.
             test_method = getattr(test_case, test_case._testMethodName)  # noqa: SLF001
             lineno = get_source_line(test_method)

python_files/vscode_pytest/__init__.py

@@ -10,7 +10,17 @@
 import pathlib
 import sys
 import traceback
-from typing import TYPE_CHECKING, Any, Dict, Generator, Literal, Protocol, TypedDict, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generator,
+    Literal,
+    NotRequired,
+    Protocol,
+    TypedDict,
+    cast,
+)
 
 import pytest
 
@@ -52,6 +62,7 @@ class TestNode(TestData):
     """A general class that handles all test data which contains children."""
 
     children: list[TestNode | TestItem | None]
+    lineno: NotRequired[str]  # Optional field for class/function nodes
 
 
 class VSCodePytestError(Exception):
@@ -830,12 +841,25 @@ def create_class_node(class_module: pytest.Class | DescribeBlock) -> TestNode:
     Keyword arguments:
     class_module -- the pytest object representing a class module.
     """
+    # Get line number for the class definition
+    class_line = ""
+    try:
+        if hasattr(class_module, "obj"):
+            import inspect
+
+            _, lineno = inspect.getsourcelines(class_module.obj)
+            class_line = str(lineno)
+    except (OSError, TypeError):
+        # If we can't get the source lines, leave lineno empty
+        pass
+
     return {
         "name": class_module.name,
         "path": get_node_path(class_module),
         "type_": "class",
         "children": [],
         "id_": get_absolute_test_id(class_module.nodeid, get_node_path(class_module)),
+        "lineno": class_line,
     }
 
 

src/client/testing/testController/common/types.ts

@@ -177,7 +177,7 @@ export interface ITestExecutionAdapter {
 }
 
 // Same types as in python_files/unittestadapter/utils.py
-export type DiscoveredTestType = 'folder' | 'file' | 'class' | 'test';
+export type DiscoveredTestType = 'folder' | 'file' | 'class' | 'function' | 'test';
 
 export type DiscoveredTestCommon = {
     path: string;
@@ -194,6 +194,7 @@ export type DiscoveredTestItem = DiscoveredTestCommon & {
 
 export type DiscoveredTestNode = DiscoveredTestCommon & {
     children: (DiscoveredTestNode | DiscoveredTestItem)[];
+    lineno?: number | string;
 };
 
 export type DiscoveredTestPayload = {

src/client/testing/testController/common/utils.ts

@@ -257,6 +257,21 @@ export function populateTestTree(
 
                     node.canResolveChildren = true;
                     node.tags = [RunTestTag, DebugTestTag];
+
+                    // Set range for class nodes (and other nodes) if lineno is available
+                    let range: Range | undefined;
+                    if ('lineno' in child && child.lineno) {
+                        if (Number(child.lineno) === 0) {
+                            range = new Range(new Position(0, 0), new Position(0, 0));
+                        } else {
+                            range = new Range(
+                                new Position(Number(child.lineno) - 1, 0),
+                                new Position(Number(child.lineno), 0),
+                            );
+                        }
+                        node.range = range;
+                    }
+
                     testRoot!.children.add(node);
                 }
                 populateTestTree(testController, child, node, resultResolver, token);