k8s.io/test-infra@v0.0.0-20240520184403-27c6b4c223d8/pkg/benchmarkjunit/main.go (about)

     1  /*
     2  Copyright 2019 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package main
    18  
    19  import (
    20  	"bytes"
    21  	"encoding/xml"
    22  	"fmt"
    23  	"io"
    24  	"os"
    25  	"os/exec"
    26  
    27  	"github.com/sirupsen/logrus"
    28  	"github.com/spf13/cobra"
    29  )
    30  
    31  type options struct {
    32  	outputFile    string
    33  	logFile       string
    34  	benchRegexp   string
    35  	extraTestArgs []string
    36  	goBinaryPath  string
    37  	passOnError   bool
    38  }
    39  
    40  func main() {
    41  	opts := &options{}
    42  	cmd := &cobra.Command{
    43  		Use:   "benchmarkjunit <packages>",
    44  		Short: "Runs go benchmarks and outputs junit xml.",
    45  		Long:  `Runs "go test -v -run='^$' -bench=. <packages>" and translates the output into JUnit XML.`,
    46  		Run: func(cmd *cobra.Command, args []string) {
    47  			run(opts, args)
    48  		},
    49  	}
    50  	cmd.Flags().StringVarP(&opts.outputFile, "output", "o", "-", "output file")
    51  	cmd.Flags().StringVarP(&opts.logFile, "log-file", "l", "", "optional output file for complete go test output. Use '-' to stream output to Stdout.")
    52  	cmd.Flags().StringVar(&opts.benchRegexp, "bench", ".", "The regexp to pass to the -bench 'go test' flag to select benchmarks to run.")
    53  	cmd.Flags().StringSliceVar(&opts.extraTestArgs, "test-arg", nil, "additional args for go test")
    54  	cmd.Flags().StringVar(&opts.goBinaryPath, "go", "go", "The location of the go binary. This flag is primarily intended for use with bazel.")
    55  	cmd.Flags().BoolVar(&opts.passOnError, "pass-on-error", false, "Indicates that benchmarkjunit should exit zero if junit is properly generated, even if benchmarks fail.")
    56  
    57  	if err := cmd.Execute(); err != nil {
    58  		logrus.WithError(err).Fatal("Command failed.")
    59  	}
    60  }
    61  
    62  func run(opts *options, args []string) {
    63  	testArgs := []string{
    64  		"test", "-v", "-run='^$'", "-bench=" + opts.benchRegexp,
    65  	}
    66  	testArgs = append(testArgs, opts.extraTestArgs...)
    67  	testArgs = append(testArgs, args...)
    68  	testCmd := exec.Command(opts.goBinaryPath, testArgs...)
    69  
    70  	logrus.Infof("Running command %q...", append([]string{opts.goBinaryPath}, testArgs...))
    71  	var testOutput []byte
    72  	var testErr error
    73  	if opts.logFile == "-" {
    74  		// Stream command output to stdout.
    75  		var buf bytes.Buffer
    76  		writer := io.MultiWriter(os.Stdout, &buf)
    77  		testCmd.Stdout = writer
    78  		testCmd.Stderr = writer
    79  		testErr = testCmd.Run()
    80  		testOutput = buf.Bytes()
    81  	} else {
    82  		testOutput, testErr = testCmd.CombinedOutput()
    83  	}
    84  	if testErr != nil {
    85  		logrus.WithError(testErr).Error("Error(s) executing benchmarks.")
    86  	}
    87  	if len(opts.logFile) > 0 && opts.logFile != "-" {
    88  		if err := os.WriteFile(opts.logFile, testOutput, 0666); err != nil {
    89  			logrus.WithError(err).Fatalf("Failed to write to log file %q.", opts.logFile)
    90  		}
    91  	}
    92  	logrus.Info("Benchmarks completed. Generating JUnit XML...")
    93  
    94  	// Now parse output to JUnit, marshal to XML, and output.
    95  	junit, err := parse(testOutput)
    96  	if err != nil {
    97  		logrus.WithField("output", string(testOutput)).WithError(err).Fatal("Error parsing 'go test' output.")
    98  	}
    99  	if len(junit.Suites) == 0 {
   100  		logrus.WithField("output", string(testOutput)).Fatal("Error: no test suites were found in the 'go test' output.")
   101  	}
   102  	junitBytes, err := xml.Marshal(junit)
   103  	if err != nil {
   104  		logrus.WithError(err).Fatal("Error marshaling parsed 'go test' output to XML.")
   105  	}
   106  	if opts.outputFile == "-" {
   107  		fmt.Println(string(junitBytes))
   108  	} else {
   109  		if err := os.WriteFile(opts.outputFile, junitBytes, 0666); err != nil {
   110  			logrus.WithError(err).Fatalf("Failed to write JUnit to output file %q.", opts.outputFile)
   111  		}
   112  	}
   113  	logrus.Info("Successfully generated JUnit XML for Benchmarks.")
   114  
   115  	if !opts.passOnError && testErr != nil {
   116  		logrus.WithError(testErr).Fatal("Exiting non-zero due to benchmark error.")
   117  	}
   118  }