gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/pkg/test/testutil/tree.go (about)

     1  // Copyright 2024 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package testutil
    16  
    17  import (
    18  	"sort"
    19  	"strings"
    20  	"testing"
    21  )
    22  
    23  // Tree represents a hierarchy of tests and sub-tests.
    24  // It is a nested structure built out of a flat list of fully-qualified
    25  // test names, and can then execute them using nested `t.Run`s.
    26  // It is useful to run a series of hierarchical Go tests in cases where
    27  // the hierarchy is not known at test compilation time.
    28  type Tree struct {
    29  	root *treeNode
    30  }
    31  
    32  // treeNode represents a hierarchy of tests and sub-tests.
    33  type treeNode struct {
    34  	// testName is only set on leaf nodes.
    35  	// It is a fully-qualified test name.
    36  	testName string
    37  
    38  	// children is only set of non-leaf nodes.
    39  	// It is a set of child nodes, mapped by their component as key.
    40  	children map[string]*treeNode
    41  }
    42  
    43  // NewTree creates a new test tree out of the given test names.
    44  // Each test name is split by `separator`, which indicates nesting.
    45  // Only leaf nodes are considered actual tests.
    46  // For example: `NewTree([]string{"a/b", "a/c", "a/c/d"}, "/")`
    47  // contains two tests: `a/b` and `a/c/d`.
    48  func NewTree(testNames []string, separator string) *Tree {
    49  	tree := &Tree{root: &treeNode{}}
    50  	for _, testName := range testNames {
    51  		n := tree.root
    52  		for _, component := range strings.Split(testName, separator) {
    53  			if component == "" {
    54  				continue
    55  			}
    56  			child, found := n.children[component]
    57  			if !found {
    58  				child = &treeNode{}
    59  				if n.children == nil {
    60  					n.children = make(map[string]*treeNode)
    61  				}
    62  				n.children[component] = child
    63  			}
    64  			n = child
    65  		}
    66  		n.testName = testName
    67  	}
    68  	return tree
    69  }
    70  
    71  // run calls `t.Run` on each test, preserving hierarchy.
    72  // `fn` is called on each leaf node with the fully-qualified test name as
    73  // argument.
    74  func (n *treeNode) run(t *testing.T, parallel bool, fn func(t *testing.T, testName string)) {
    75  	t.Helper()
    76  	if len(n.children) == 0 { // Leaf node.
    77  		fn(t, n.testName)
    78  		return
    79  	}
    80  	childNames := make([]string, 0, len(n.children))
    81  	for childName := range n.children {
    82  		childNames = append(childNames, childName)
    83  	}
    84  	sort.Strings(childNames)
    85  	for _, childName := range childNames {
    86  		childNode := n.children[childName]
    87  		t.Run(childName, func(t *testing.T) {
    88  			if parallel {
    89  				t.Parallel()
    90  			}
    91  			childNode.run(t, parallel, fn)
    92  		})
    93  	}
    94  }
    95  
    96  // Run calls `t.Run` on each leaf test, preserving test hierarchy.
    97  // `fn` is called on each leaf node with the fully-qualified test name as
    98  // argument.
    99  func (tree *Tree) Run(t *testing.T, fn func(t *testing.T, testName string)) {
   100  	t.Helper()
   101  	tree.root.run(t, false, fn)
   102  }
   103  
   104  // RunParallel calls `t.Run` on each test in parallel, preserving hierarchy.
   105  // `fn` is called on each leaf node with the fully-qualified test name as
   106  // argument.
   107  // `fn` does not need to call `t.Parallel`.
   108  func (tree *Tree) RunParallel(t *testing.T, fn func(t *testing.T, testName string)) {
   109  	t.Helper()
   110  	tree.root.run(t, true, fn)
   111  }