-
Notifications
You must be signed in to change notification settings - Fork 567
/
main.go
179 lines (166 loc) · 5.42 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
package main
import (
"bufio"
"bytes"
"context"
"encoding/json"
"flag"
"fmt"
"io"
"os"
"os/exec"
"strings"
"syscall"
"time"
"github.com/pachyderm/pachyderm/v2/src/internal/errors"
"github.com/pachyderm/pachyderm/v2/src/internal/log"
"github.com/pachyderm/pachyderm/v2/src/internal/pctx"
"github.com/pachyderm/pachyderm/v2/src/internal/proc"
"go.uber.org/zap"
"k8s.io/kubectl/pkg/util/slice"
)
type testOutput struct {
Time time.Time
Action string
Package string
Output string
}
func main() {
log.InitPachctlLogger()
ctx := pctx.Background("test-collector")
tags := flag.String("tags", "", "Tags to run, for example k8s. Tests without this flag will not be selected.")
exclusiveTags := flag.Bool("exclusiveTags", true, "If true, ONLY tests with the specified tags will run. If false, "+
"the default behavior of 'go test' of including tagged and untagged tests is used.")
fileName := flag.String("file", "tests_to_run.csv", "Output file listing the packages and tests to run. Used by the runner script.")
pkg := flag.String("pkg", "./...", "Package to run defaults to all packages.")
threadPool := flag.Int("procs", 0, "GOMAXPROCS value for the go test -list sbcommand.")
flag.Parse()
err := run(ctx, *tags, *exclusiveTags, *fileName, *pkg, *threadPool)
if err != nil {
log.Exit(ctx, "Error during tests splitting", zap.Error(err))
}
os.Exit(0)
}
func run(ctx context.Context, tags string, exclusiveTags bool, fileName string, pkg string, threadPool int) error {
var tagsArg string
if tags != "" {
tagsArg = fmt.Sprintf("-tags=%s", tags)
}
testIdsTagged, err := testNames(ctx, pkg, threadPool, tagsArg)
if err != nil {
return errors.EnsureStack(err)
}
var testIds map[string][]string
if exclusiveTags && tags != "" {
// set difference to get ONLY tagged tests
var err error
testIdsUntagged, err := testNames(ctx, pkg, threadPool, "") // collect for set difference
if err != nil {
return errors.EnsureStack(err)
}
testIds = subtractTestSet(testIdsTagged, testIdsUntagged)
} else {
testIds = testIdsTagged
}
log.Info(ctx, "tests and packages collected", zap.Any("tests", testIds))
err = outputToFile(fileName, testIds)
return err
}
// get tests that are tagged, but not in the untagged list since that is inclusive
func subtractTestSet(testIdsTagged map[string][]string, testIdsUntagged map[string][]string) map[string][]string {
resultTests := map[string][]string{}
for pkg, testsTagged := range testIdsTagged {
testsUntagged, ok := testIdsUntagged[pkg]
for _, testNameTagged := range testsTagged {
if !ok || !slice.ContainsString(testsUntagged, testNameTagged, func(s string) string { return s }) {
if _, ok := resultTests[pkg]; !ok {
resultTests[pkg] = []string{testNameTagged}
} else {
resultTests[pkg] = append(resultTests[pkg], testNameTagged)
}
}
}
}
return resultTests
}
func testNames(ctx context.Context, pkg string, threadPool int, addtlCmdArgs ...string) (map[string][]string, error) {
findTestArgs := append([]string{"test", pkg, "-json", "-list=."}, addtlCmdArgs...)
cmd := exec.Command("go", findTestArgs...)
stdout, err := cmd.StdoutPipe()
if err != nil {
return nil, errors.EnsureStack(err)
}
cmd.Stderr = log.WriterAt(log.ChildLogger(ctx, "stderr"), log.InfoLevel)
if threadPool > 0 {
cmd.Env = os.Environ()
cmd.Env = append(cmd.Env, fmt.Sprintf("GOMAXPROCS=%d", threadPool)) // This prevents the command from running wild eating up processes in the pipelines
}
cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
monitorCtx, monitorCancel := context.WithCancel(pctx.Child(ctx, "monitor go process"))
defer monitorCancel()
go proc.MonitorProcessGroup(monitorCtx, cmd.SysProcAttr.Pgid)
err = cmd.Start()
if err != nil {
return nil, errors.EnsureStack(err)
}
testNames, err := readTests(stdout)
if err != nil {
return nil, err
}
err = cmd.Wait()
if err != nil {
return nil, errors.EnsureStack(err)
}
return testNames, nil
}
func readTests(stdout io.Reader) (map[string][]string, error) {
var testNames = map[string][]string{}
scanner := bufio.NewScanner(stdout)
for scanner.Scan() {
testInfo := &testOutput{}
raw := scanner.Bytes()
if !bytes.HasPrefix(raw, []byte("{")) { // dependency download junk got shared
continue
}
if err := json.Unmarshal(raw, testInfo); err != nil {
return nil, errors.Wrapf(err, "parsing json: %s", string(raw))
}
if testInfo.Action == "output" {
output := strings.Trim(testInfo.Output, "\n ")
if output != "" && !strings.HasPrefix(output, "Benchmark") &&
!strings.HasPrefix(output, "ExampleAPIClient_") &&
!strings.HasPrefix(output, "? ") &&
!strings.HasPrefix(output, "ok ") {
if _, ok := testNames[testInfo.Package]; !ok {
testNames[testInfo.Package] = []string{output}
} else {
testNames[testInfo.Package] = append(testNames[testInfo.Package], output)
}
}
}
}
return testNames, nil
}
func outputToFile(fileName string, pkgTests map[string][]string) error {
tmpFileName := fmt.Sprintf("%s.tmp", fileName)
f, err := os.Create(tmpFileName)
if err != nil {
return errors.EnsureStack(err)
}
defer f.Close()
w := bufio.NewWriter(f)
for pkg, tests := range pkgTests {
for _, test := range tests {
_, err := w.WriteString(fmt.Sprintf("%s,%s\n", pkg, test))
if err != nil {
return errors.EnsureStack(err)
}
}
}
err = errors.EnsureStack(w.Flush())
if err != nil {
return err
}
err = os.Rename(tmpFileName, fileName)
return errors.EnsureStack(err)
}