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

# Licensed under the MIT License.

import argparse

import enum

import inspect

import os

import pathlib

import sys

import unittest

from typing import List, Tuple, Union

script_dir = pathlib.Path(__file__).parent.parent

sys.path.append(os.fspath(script_dir))

sys.path.append(os.fspath(script_dir / "lib" / "python"))

from typing_extensions import TypedDict

# 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

id_: str

class TestItem(TestData):

lineno: str

runID: 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."""

## figure out if we are folder, file, or class

id_gen = path

if type_ == TestNodeTypeEnum.folder or type_ == TestNodeTypeEnum.file:

id_gen = path

else:

# means we have to build test node for class

id_gen = path + "\\" + name

return {"path": path, "name": name, "type_": type_, "children": [], "id_": id_gen}

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 following format in line 196>,

}

],

"id_": <class path path following format after path>

}

],

"id_": <file path>

}

],

"id_": <test_directory path>

}

"""

error = []

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"):

error.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 = {

"name": function_name,

"path": file_path,

"lineno": lineno,

"type_": TestNodeTypeEnum.test,

"id_": file_path + "\\" + class_name + "\\" + function_name,

"runID": test_id,

} # concatenate class name and function test name

current_node["children"].append(test_node)

if not root["children"]:

root = None

return root, error

def parse_unittest_args(args: List[str]) -> Tuple[str, str, Union[str, None]]:

"""Parse command-line arguments that should be forwarded to unittest to perform discovery.

Valid unittest arguments are: -v, -s, -p, -t and their long-form counterparts,

however we only care about the last three.

The returned tuple contains the following items

- start_directory: The directory where to start discovery, defaults to .

- pattern: The pattern to match test files, defaults to test*.py

- top_level_directory: The top-level directory of the project, defaults to None, and unittest will use start_directory behind the scenes.

"""

arg_parser = argparse.ArgumentParser()

arg_parser.add_argument("--start-directory", "-s", default=".")

arg_parser.add_argument("--pattern", "-p", default="test*.py")

arg_parser.add_argument("--top-level-directory", "-t", default=None)

parsed_args, _ = arg_parser.parse_known_args(args)

return (

parsed_args.start_directory,

parsed_args.pattern,

parsed_args.top_level_directory,

)