github.com/mwhudson/juju@v0.0.0-20160512215208-90ff01f3497f/testing/cmd.go (about)

     1  // Copyright 2013 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package testing
     5  
     6  import (
     7  	"bytes"
     8  	"fmt"
     9  	"io/ioutil"
    10  	"strings"
    11  
    12  	"github.com/juju/cmd"
    13  	jc "github.com/juju/testing/checkers"
    14  	gc "gopkg.in/check.v1"
    15  	"launchpad.net/gnuflag"
    16  )
    17  
    18  // NewFlagSet creates a new flag set using the standard options, particularly
    19  // the option to stop the gnuflag methods from writing to StdErr or StdOut.
    20  func NewFlagSet() *gnuflag.FlagSet {
    21  	fs := gnuflag.NewFlagSet("", gnuflag.ContinueOnError)
    22  	fs.SetOutput(ioutil.Discard)
    23  	return fs
    24  }
    25  
    26  // InitCommand will create a new flag set, and call the Command's SetFlags and
    27  // Init methods with the appropriate args.
    28  func InitCommand(c cmd.Command, args []string) error {
    29  	f := NewFlagSet()
    30  	c.SetFlags(f)
    31  	if err := f.Parse(c.AllowInterspersedFlags(), args); err != nil {
    32  		return err
    33  	}
    34  	return c.Init(f.Args())
    35  }
    36  
    37  // Context creates a simple command execution context with the current
    38  // dir set to a newly created directory within the test directory.
    39  func Context(c *gc.C) *cmd.Context {
    40  	return &cmd.Context{
    41  		Dir:    c.MkDir(),
    42  		Stdin:  &bytes.Buffer{},
    43  		Stdout: &bytes.Buffer{},
    44  		Stderr: &bytes.Buffer{},
    45  	}
    46  }
    47  
    48  // ContextForDir creates a simple command execution context with the current
    49  // dir set to the specified directory.
    50  func ContextForDir(c *gc.C, dir string) *cmd.Context {
    51  	return &cmd.Context{
    52  		Dir:    dir,
    53  		Stdin:  &bytes.Buffer{},
    54  		Stdout: &bytes.Buffer{},
    55  		Stderr: &bytes.Buffer{},
    56  	}
    57  }
    58  
    59  // Stdout takes a command Context that we assume has been created in this
    60  // package, and gets the content of the Stdout buffer as a string.
    61  func Stdout(ctx *cmd.Context) string {
    62  	buff, ok := ctx.Stdout.(*bytes.Buffer)
    63  	if ok {
    64  		return buff.String()
    65  	}
    66  	return "not a buffer"
    67  }
    68  
    69  // Stderr takes a command Context that we assume has been created in this
    70  // package, and gets the content of the Stderr buffer as a string.
    71  func Stderr(ctx *cmd.Context) string {
    72  	return ctx.Stderr.(*bytes.Buffer).String()
    73  }
    74  
    75  // RunCommand runs a command with the specified args.  The returned error
    76  // may come from either the parsing of the args, the command initialisation, or
    77  // the actual running of the command.  Access to the resulting output streams
    78  // is provided through the returned context instance.
    79  func RunCommand(c *gc.C, com cmd.Command, args ...string) (*cmd.Context, error) {
    80  	var context = Context(c)
    81  	if err := InitCommand(com, args); err != nil {
    82  		if err != nil {
    83  			fmt.Fprintf(context.Stderr, "error: %v\n", err)
    84  		}
    85  		return context, err
    86  	}
    87  	return context, com.Run(context)
    88  }
    89  
    90  // RunCommandInDir works like RunCommand, but runs with a context that uses dir.
    91  func RunCommandInDir(c *gc.C, com cmd.Command, args []string, dir string) (*cmd.Context, error) {
    92  	if err := InitCommand(com, args); err != nil {
    93  		return nil, err
    94  	}
    95  	var context = ContextForDir(c, dir)
    96  	return context, com.Run(context)
    97  }
    98  
    99  // TestInit checks that a command initialises correctly with the given set of
   100  // arguments.
   101  func TestInit(c *gc.C, com cmd.Command, args []string, errPat string) {
   102  	err := InitCommand(com, args)
   103  	if errPat != "" {
   104  		c.Assert(err, gc.ErrorMatches, errPat)
   105  	} else {
   106  		c.Assert(err, jc.ErrorIsNil)
   107  	}
   108  }
   109  
   110  // ExtractCommandsFromHelpOutput takes the standard output from the
   111  // command context and looks for the "commands:" string and returns the
   112  // commands output after that.
   113  func ExtractCommandsFromHelpOutput(ctx *cmd.Context) []string {
   114  	var namesFound []string
   115  	commandHelp := strings.SplitAfter(Stdout(ctx), "commands:")[1]
   116  	commandHelp = strings.TrimSpace(commandHelp)
   117  	for _, line := range strings.Split(commandHelp, "\n") {
   118  		namesFound = append(namesFound, strings.TrimSpace(strings.Split(line, " - ")[0]))
   119  	}
   120  	return namesFound
   121  }