# Copyright (c) Microsoft Corporation. All rights reserved.

# Licensed under the MIT License.

import enum

import inspect

import os

import pathlib

import unittest

from typing import List, Tuple, TypedDict, Union

# Types

# Inherit from str so it's JSON serializable.

class TestNodeTypeEnum(str, enum.Enum):

class_ = "class"

file = "file"

folder = "folder"

test = "test"

class TestData(TypedDict):

name: str

path: str

type_: TestNodeTypeEnum

class TestItem(TestData):

id_: str

lineno: str

class TestNode(TestData):

children: "List[TestNode | TestItem]"

# Helper functions for data retrieval.

def get_test_case(suite):

"""Iterate through a unittest test suite and return all test cases."""

for test in suite:

if isinstance(test, unittest.TestCase):

yield test

else:

for test_case in get_test_case(test):

yield test_case

def get_source_line(obj) -> str:

"""Get the line number of a test case start line."""

try:

sourcelines, lineno = inspect.getsourcelines(obj)

except:

try:

# tornado-specific, see https://github.com/microsoft/vscode-python/issues/17285.

sourcelines, lineno = inspect.getsourcelines(obj.orig_method)

except:

return "*"

# Return the line number of the first line of the test case definition.

for i, v in enumerate(sourcelines):

if v.strip().startswith(("def", "async def")):

return str(lineno + i)

return "*"

# Helper functions for test tree building.

def build_test_node(path: str, name: str, type_: TestNodeTypeEnum) -> TestNode:

"""Build a test node with no children. A test node can be a folder, a file or a class."""

return {

"path": path,

"name": name,

"type_": type_,

"children": [],

}

def get_child_node(

name: str, path: str, type_: TestNodeTypeEnum, root: TestNode

) -> TestNode:

"""Find a child node in a test tree given its name and type. If the node doesn't exist, create it."""

try:

result = next(

node

for node in root["children"]

if node["name"] == name and node["type_"] == type_

)

except StopIteration:

result = build_test_node(path, name, type_)

root["children"].append(result)

return result # type:ignore

def build_test_tree(

suite: unittest.TestSuite, test_directory: str

) -> Tuple[Union[TestNode, None], List[str]]:

"""Build a test tree from a unittest test suite.

This function returns the test tree, and any errors found by unittest.

If no tests were discovered, return `None` and a list of errors (if any).

Test tree structure:

{

"path": <test_directory path>,

"type": "folder",

"name": <folder name>,

"children": [

{ files and folders }

...

{

"path": <file path>,

"name": filename.py,

"type_": "file",

"children": [

{

"path": <class path>,

"name": <class name>,

"type_": "class",

"children": [

{

"path": <test path>,

"name": <test name>,

"type_": "test",

"lineno": <line number>

"id": <test case id>,

}

]

}

]

}

]

}

"""

errors = []

directory_path = pathlib.PurePath(test_directory)

root = build_test_node(test_directory, directory_path.name, TestNodeTypeEnum.folder)

for test_case in get_test_case(suite):

test_id = test_case.id()

if test_id.startswith("unittest.loader._FailedTest"):

errors.append(str(test_case._exception)) # type: ignore

else:

# Get the static test path components: filename, class name and function name.

components = test_id.split(".")

*folders, filename, class_name, function_name = components

py_filename = f"{filename}.py"

current_node = root

# Find/build nodes for the intermediate folders in the test path.

for folder in folders:

current_node = get_child_node(

folder,

os.fsdecode(pathlib.PurePath(current_node["path"], folder)),

TestNodeTypeEnum.folder,

current_node,

)

# Find/build file node.

path_components = [test_directory] + folders + [py_filename]

file_path = os.fsdecode(pathlib.PurePath("/".join(path_components)))

current_node = get_child_node(

py_filename, file_path, TestNodeTypeEnum.file, current_node

)

# Find/build class node.

current_node = get_child_node(

class_name, file_path, TestNodeTypeEnum.class_, current_node

)

# Get test line number.

test_method = getattr(test_case, test_case._testMethodName)

lineno = get_source_line(test_method)

# Add test node.

test_node: TestItem = {

"id_": test_id,

"name": function_name,

"path": file_path,

"lineno": lineno,

"type_": TestNodeTypeEnum.test,

}

current_node["children"].append(test_node)

if not root["children"]:

root = None

return root, errors