go.uber.org/yarpc@v1.72.1/internal/service-test/main.go (about)

     1  // Copyright (c) 2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package main
    22  
    23  import (
    24  	"errors"
    25  	"flag"
    26  	"fmt"
    27  	"log"
    28  	"os"
    29  	"os/signal"
    30  	"path/filepath"
    31  	"time"
    32  )
    33  
    34  var (
    35  	flagSet            = flag.NewFlagSet("service-test", flag.ExitOnError)
    36  	flagDir            = flagSet.String("dir", "", "The relative directory to operate from, defaults to current directory")
    37  	flagConfigFilePath = flagSet.String("file", "service-test.yaml", "The configuration file to use relative to the context directory")
    38  	flagTimeout        = flagSet.Duration("timeout", 5*time.Second, "The time to wait until timing out")
    39  	flagNoVerifyOutput = flagSet.Bool("no-validate-output", false, "Do not validate output and just run the commands")
    40  	flagDebug          = flagSet.Bool("debug", false, "Log debug information")
    41  
    42  	errSignal = errors.New("signal")
    43  )
    44  
    45  func main() {
    46  	if err := flagSet.Parse(os.Args[1:]); err != nil {
    47  		log.Fatal(err)
    48  	}
    49  	if err := do(
    50  		*flagDir,
    51  		*flagConfigFilePath,
    52  		*flagTimeout,
    53  		!(*flagNoVerifyOutput),
    54  		*flagDebug,
    55  	); err != nil {
    56  		log.Fatal(err)
    57  	}
    58  }
    59  
    60  func do(
    61  	dir string,
    62  	configFilePath string,
    63  	timeout time.Duration,
    64  	validateOutput bool,
    65  	debug bool,
    66  ) (err error) {
    67  	cmds, err := newCmds(configFilePath, dir, debug)
    68  	if err != nil {
    69  		return err
    70  	}
    71  	defer cleanupCmds(cmds, validateOutput, err)
    72  	signalC := make(chan os.Signal, 1)
    73  	signal.Notify(signalC, os.Interrupt)
    74  	go func() {
    75  		for range signalC {
    76  			cleanupCmds(cmds, validateOutput, errSignal)
    77  		}
    78  		os.Exit(1)
    79  	}()
    80  
    81  	errC := make(chan error)
    82  	go func() { errC <- runCmds(cmds) }()
    83  	select {
    84  	case err := <-errC:
    85  		if err != nil {
    86  			return err
    87  		}
    88  	case <-time.After(timeout):
    89  		return fmt.Errorf("timed out after %v", timeout)
    90  	}
    91  	if validateOutput {
    92  		return validateCmds(cmds)
    93  	}
    94  	return nil
    95  }
    96  
    97  func newCmds(configFilePath string, dir string, debug bool) ([]*cmd, error) {
    98  	config, err := newConfig(filepath.Join(dir, configFilePath))
    99  	if err != nil {
   100  		return nil, err
   101  	}
   102  	return config.Cmds(dir, debug)
   103  }
   104  
   105  func runCmds(cmds []*cmd) error {
   106  	for i := 0; i < len(cmds)-1; i++ {
   107  		cmd := cmds[i]
   108  		if err := cmd.Start(); err != nil {
   109  			return err
   110  		}
   111  		defer func() {
   112  			cmd.Kill()
   113  			_ = cmd.Wait()
   114  		}()
   115  	}
   116  	cmd := cmds[len(cmds)-1]
   117  	if err := cmd.Start(); err != nil {
   118  		return err
   119  	}
   120  	return cmd.Wait()
   121  }
   122  
   123  func validateCmds(cmds []*cmd) error {
   124  	for _, cmd := range cmds {
   125  		if err := cmd.Validate(); err != nil {
   126  			return err
   127  		}
   128  	}
   129  	return nil
   130  }
   131  
   132  func cleanupCmds(cmds []*cmd, validateOutput bool, err error) {
   133  	for _, cmd := range cmds {
   134  		cmd.Clean(validateOutput && err == nil)
   135  	}
   136  }