Refactor test processing and folder construction in vscode_pytest module

Author: eleanorjboyd

python_files/vscode_pytest/__init__.py

@@ -10,7 +10,7 @@
 import pathlib
 import sys
 import traceback
-from typing import TYPE_CHECKING, Any, Dict, Generator, Literal, TypedDict
+from typing import TYPE_CHECKING, Any, Dict, Generator, Literal, Protocol, TypedDict
 
 import pytest
 
@@ -25,6 +25,13 @@
     USES_PYTEST_DESCRIBE = True
 
 
+class HasPathOrFspath(Protocol):
+    """Protocol defining objects that have either a path or fspath attribute."""
+
+    path: pathlib.Path | None = None
+    fspath: Any | None = None
+
+
 class TestData(TypedDict):
     """A general class that all test objects inherit from."""
 
@@ -522,11 +529,131 @@ def pytest_sessionfinish(session, exitstatus):
         send_message(payload)
 
 
+def construct_nested_folders(
+    file_nodes_dict: dict[str, TestNode],
+    session_node: TestNode,  # session_node['path'] is a pathlib.Path object
+    session_children_dict: dict[str, TestNode]
+) -> dict[str, TestNode]:
+    """Iterate through all files and construct them into nested folders.
+
+    Keyword arguments:
+    file_nodes_dict -- Dictionary of all file nodes
+    session_node -- The session node that will be parent to the folder structure
+    session_children_dict -- Dictionary of session's children nodes indexed by ID
+
+    Returns:
+    dict[str, TestNode] -- Updated session_children_dict with folder nodes added
+    """
+    created_files_folders_dict: dict[str, TestNode] = {}
+    for file_node in file_nodes_dict.values():
+        # Iterate through all the files that exist and construct them into nested folders.
+        root_folder_node: TestNode
+        try:
+            root_folder_node: TestNode = build_nested_folders(
+                file_node, created_files_folders_dict, session_node
+            )
+        except ValueError:
+            # This exception is raised when the session node is not a parent of the file node.
+            print(
+                "[vscode-pytest]: Session path not a parent of test paths, adjusting session node to common parent."
+            )
+            # IMPORTANT: Use session_node["path"] directly as it's already a pathlib.Path object
+            # Do NOT use get_node_path(session_node["path"]) as get_node_path expects pytest objects,
+            # not Path objects directly.
+            common_parent = os.path.commonpath([file_node["path"], session_node["path"]])
+            common_parent_path = pathlib.Path(common_parent)
+            print("[vscode-pytest]: Session node now set to: ", common_parent)
+            session_node["path"] = common_parent_path  # pathlib.Path
+            session_node["id_"] = common_parent  # str
+            session_node["name"] = common_parent_path.name  # str
+            root_folder_node = build_nested_folders(
+                file_node, created_files_folders_dict, session_node
+            )
+        # The final folder we get to is the highest folder in the path
+        # and therefore we add this as a child to the session.
+        root_id = root_folder_node.get("id_")
+        if root_id and root_id not in session_children_dict:
+            session_children_dict[root_id] = root_folder_node
+
+    return session_children_dict
+
+
+def process_parameterized_test(
+    test_case: pytest.Item,  # Must have callspec attribute (parameterized test)
+    test_node: TestItem,
+    function_nodes_dict: dict[str, TestNode],
+    file_nodes_dict: dict[str, TestNode]
+) -> TestNode:
+    """Process a parameterized test case and create appropriate function nodes.
+
+    Keyword arguments:
+    test_case -- the parameterized pytest test case
+    test_node -- the test node created from the test case
+    function_nodes_dict -- dictionary of function nodes indexed by ID
+    file_nodes_dict -- dictionary of file nodes indexed by path
+
+    Returns:
+    TestNode -- the node to use for further processing (function node or original test node)
+    """
+    function_name: str = ""
+    # parameterized test cases cut the repetitive part of the name off.
+    parent_part, parameterized_section = test_node["name"].split("[", 1)
+    test_node["name"] = "[" + parameterized_section
+
+    first_split = test_case.nodeid.rsplit(
+        "::", 1
+    )  # splits the parameterized test name from the rest of the nodeid
+    second_split = first_split[0].rsplit(
+        ".py", 1
+    )  # splits the file path from the rest of the nodeid
+
+    class_and_method = second_split[1] + "::"  # This has "::" separator at both ends
+    # construct the parent id, so it is absolute path :: any class and method :: parent_part
+    parent_id = os.fspath(get_node_path(test_case)) + class_and_method + parent_part
+
+    try:
+        function_name = test_case.originalname  # type: ignore
+        function_test_node = function_nodes_dict[parent_id]
+    except AttributeError:  # actual error has occurred
+        ERRORS.append(
+            f"unable to find original name for {test_case.name} with parameterization detected."
+        )
+        raise VSCodePytestError(
+            "Unable to find original name for parameterized test case"
+        ) from None
+    except KeyError:
+        function_test_node: TestNode = create_parameterized_function_node(
+            function_name, get_node_path(test_case), parent_id
+        )
+        function_nodes_dict[parent_id] = function_test_node
+
+    if test_node not in function_test_node["children"]:
+        function_test_node["children"].append(test_node)
+
+    # Check if the parent node of the function is file, if so create/add to this file node.
+    if isinstance(test_case.parent, pytest.File):
+        # calculate the parent path of the test case
+        parent_path = get_node_path(test_case.parent)
+        try:
+            parent_test_case = file_nodes_dict[os.fspath(parent_path)]
+        except KeyError:
+            parent_test_case = create_file_node(parent_path)
+            file_nodes_dict[os.fspath(parent_path)] = parent_test_case
+        if function_test_node not in parent_test_case["children"]:
+            parent_test_case["children"].append(function_test_node)
+
+    # Return the function node as the test node to handle subsequent nesting
+    return function_test_node
+
+
 def build_test_tree(session: pytest.Session) -> TestNode:
     """Builds a tree made up of testing nodes from the pytest session.
 
     Keyword arguments:
-    session -- the pytest session object.
+    session -- the pytest session object that contains test items.
+
+    Returns:
+    TestNode -- The root node of the constructed test tree.
     """
     session_node = create_session_node(session)
     session_children_dict: dict[str, TestNode] = {}
@@ -542,54 +669,8 @@ def build_test_tree(session: pytest.Session) -> TestNode:
     for test_case in session.items:
         test_node = create_test_node(test_case)
         if hasattr(test_case, "callspec"):  # This means it is a parameterized test.
-            function_name: str = ""
-            # parameterized test cases cut the repetitive part of the name off.
-            parent_part, parameterized_section = test_node["name"].split("[", 1)
-            test_node["name"] = "[" + parameterized_section
-
-            first_split = test_case.nodeid.rsplit(
-                "::", 1
-            )  # splits the parameterized test name from the rest of the nodeid
-            second_split = first_split[0].rsplit(
-                ".py", 1
-            )  # splits the file path from the rest of the nodeid
-
-            class_and_method = second_split[1] + "::"  # This has "::" separator at both ends
-            # construct the parent id, so it is absolute path :: any class and method :: parent_part
-            parent_id = os.fspath(get_node_path(test_case)) + class_and_method + parent_part
-            # file, middle, param = test_case.nodeid.rsplit("::", 2)
-            # parent_id = test_case.nodeid.rsplit("::", 1)[0] + "::" + parent_part
-            # parent_path = os.fspath(get_node_path(test_case)) + "::" + parent_part
-            try:
-                function_name = test_case.originalname  # type: ignore
-                function_test_node = function_nodes_dict[parent_id]
-            except AttributeError:  # actual error has occurred
-                ERRORS.append(
-                    f"unable to find original name for {test_case.name} with parameterization detected."
-                )
-                raise VSCodePytestError(
-                    "Unable to find original name for parameterized test case"
-                ) from None
-            except KeyError:
-                function_test_node: TestNode = create_parameterized_function_node(
-                    function_name, get_node_path(test_case), parent_id
-                )
-                function_nodes_dict[parent_id] = function_test_node
-            if test_node not in function_test_node["children"]:
-                function_test_node["children"].append(test_node)
-            # Check if the parent node of the function is file, if so create/add to this file node.
-            if isinstance(test_case.parent, pytest.File):
-                # calculate the parent path of the test case
-                parent_path = get_node_path(test_case.parent)
-                try:
-                    parent_test_case = file_nodes_dict[os.fspath(parent_path)]
-                except KeyError:
-                    parent_test_case = create_file_node(parent_path)
-                    file_nodes_dict[os.fspath(parent_path)] = parent_test_case
-                if function_test_node not in parent_test_case["children"]:
-                    parent_test_case["children"].append(function_test_node)
-            # If the parent is not a file, it is a class, add the function node as the test node to handle subsequent nesting.
-            test_node = function_test_node
+            # Process parameterized test and get the function node to use for further processing
+            test_node = process_parameterized_test(test_case, test_node, function_nodes_dict, file_nodes_dict)
         if isinstance(test_case.parent, pytest.Class) or (
             USES_PYTEST_DESCRIBE and isinstance(test_case.parent, DescribeBlock)
         ):
@@ -636,41 +717,16 @@ def build_test_tree(session: pytest.Session) -> TestNode:
                 parent_test_case = create_file_node(parent_path)
                 file_nodes_dict[os.fspath(parent_path)] = parent_test_case
             parent_test_case["children"].append(test_node)
-    created_files_folders_dict: dict[str, TestNode] = {}
-    for file_node in file_nodes_dict.values():
-        # Iterate through all the files that exist and construct them into nested folders.
-        root_folder_node: TestNode
-        try:
-            root_folder_node: TestNode = build_nested_folders(
-                file_node, created_files_folders_dict, session_node
-            )
-        except ValueError:
-            # This exception is raised when the session node is not a parent of the file node.
-            print(
-                "[vscode-pytest]: Session path not a parent of test paths, adjusting session node to common parent."
-            )
-            common_parent = os.path.commonpath([file_node["path"], get_node_path(session)])
-            common_parent_path = pathlib.Path(common_parent)
-            print("[vscode-pytest]: Session node now set to: ", common_parent)
-            session_node["path"] = common_parent_path  # pathlib.Path
-            session_node["id_"] = common_parent  # str
-            session_node["name"] = common_parent_path.name  # str
-            root_folder_node = build_nested_folders(
-                file_node, created_files_folders_dict, session_node
-            )
-        # The final folder we get to is the highest folder in the path
-        # and therefore we add this as a child to the session.
-        root_id = root_folder_node.get("id_")
-        if root_id and root_id not in session_children_dict:
-            session_children_dict[root_id] = root_folder_node
+    # Process all files and construct them into nested folders
+    session_children_dict = construct_nested_folders(file_nodes_dict, session_node, session_children_dict)
     session_node["children"] = list(session_children_dict.values())
     return session_node
 
 
 def build_nested_folders(
-    file_node: TestNode,
-    created_files_folders_dict: dict[str, TestNode],
-    session_node: TestNode,
+    file_node: TestNode,  # A file node to build folders for
+    created_files_folders_dict: dict[str, TestNode],  # Cache of created folders indexed by path
+    session_node: TestNode,  # The session node containing path information
 ) -> TestNode:
     """Takes a file or folder and builds the nested folder structure for it.
 
@@ -851,10 +907,17 @@ class CoveragePayloadDict(Dict):
     error: str | None  # Currently unused need to check
 
 
-def get_node_path(node: Any) -> pathlib.Path:
+def get_node_path(node: pytest.Session | pytest.Item | pytest.File | pytest.Class | pytest.Module | HasPathOrFspath) -> pathlib.Path:
     """A function that returns the path of a node given the switch to pathlib.Path.
 
     It also evaluates if the node is a symlink and returns the equivalent path.
+
+    Parameters:
+        node: A pytest object or any object that has a path or fspath attribute.
+            Do NOT pass a pathlib.Path object directly; use it directly instead.
+
+    Returns:
+        pathlib.Path: The resolved path for the node.
     """
     node_path = getattr(node, "path", None) or pathlib.Path(node.fspath)