// Copyright © 2023 Ory Corp // SPDX-License-Identifier: Apache-2.0 package cmdx import ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "net/url" "os" "testing" "golang.org/x/sync/errgroup" "github.com/spf13/cobra" "github.com/spf13/pflag" "github.com/stretchr/testify/require" "github.com/pkg/errors" ) var ( // ErrNoPrintButFail is returned to detect a failure state that was already reported to the user in some way ErrNoPrintButFail = fmt.Errorf("this error should never be printed") debugStdout, debugStderr = io.Discard, io.Discard ) func init() { if os.Getenv("DEBUG") != "" { debugStdout = os.Stdout debugStderr = os.Stderr } } // FailSilently is supposed to be used within a commands RunE function. // It silences cobras error handling and returns the ErrNoPrintButFail error. func FailSilently(cmd *cobra.Command) error { cmd.SilenceErrors = true cmd.SilenceUsage = true return errors.WithStack(ErrNoPrintButFail) } // Must fatals with the optional message if err is not nil. // Deprecated: do not use this function in commands, as it makes it impossible to test them. Instead, return the error. func Must(err error, message string, args ...interface{}) { if err == nil { return } _, _ = fmt.Fprintf(os.Stderr, message+"\n", args...) os.Exit(1) } // CheckResponse fatals if err is nil or the response.StatusCode does not match the expectedStatusCode // Deprecated: do not use this function in commands, as it makes it impossible to test them. Instead, return the error. func CheckResponse(err error, expectedStatusCode int, response *http.Response) { Must(err, "Command failed because error occurred: %s", err) if response.StatusCode != expectedStatusCode { out, err := io.ReadAll(response.Body) if err != nil { out = []byte{} } pretty, err := json.MarshalIndent(json.RawMessage(out), "", "\t") if err == nil { out = pretty } Fatalf( `Command failed because status code %d was expected but code %d was received. Response payload: %s`, expectedStatusCode, response.StatusCode, out, ) } } // FormatResponse takes an object and prints a json.MarshalIdent version of it or fatals. // Deprecated: do not use this function in commands, as it makes it impossible to test them. Instead, return the error. func FormatResponse(o interface{}) string { out, err := json.MarshalIndent(o, "", "\t") Must(err, `Command failed because an error occurred while prettifying output: %s`, err) return string(out) } // Fatalf prints to os.Stderr and exists with code 1. // Deprecated: do not use this function in commands, as it makes it impossible to test them. Instead, return the error. func Fatalf(message string, args ...interface{}) { if len(args) > 0 { _, _ = fmt.Fprintf(os.Stderr, message+"\n", args...) } else { _, _ = fmt.Fprintln(os.Stderr, message) } os.Exit(1) } // CallbackWriter will execute each callback once the message is received. // The full matched message is passed to the callback. An error returned from the callback is returned by Write. type CallbackWriter struct { Callbacks map[string]func([]byte) error buf bytes.Buffer } func (c *CallbackWriter) Write(msg []byte) (int, error) { for p, cb := range c.Callbacks { if bytes.Contains(msg, []byte(p)) { if err := cb(msg); err != nil { return 0, err } } } return c.buf.Write(msg) } func (c *CallbackWriter) String() string { return c.buf.String() } var _ io.Writer = (*CallbackWriter)(nil) func prepareCmd(cmd *cobra.Command, stdIn io.Reader, stdOut, stdErr io.Writer, args []string) { cmd.SetIn(stdIn) outs := []io.Writer{debugStdout} if stdOut != nil { outs = append(outs, stdOut) } cmd.SetOut(io.MultiWriter(outs...)) errs := []io.Writer{debugStderr} if stdErr != nil { errs = append(errs, stdErr) } cmd.SetErr(io.MultiWriter(errs...)) if args == nil { args = []string{} } cmd.SetArgs(args) } // ExecBackgroundCtx runs the cobra command in the background. func ExecBackgroundCtx(ctx context.Context, cmd *cobra.Command, stdIn io.Reader, stdOut, stdErr io.Writer, args ...string) *errgroup.Group { prepareCmd(cmd, stdIn, stdOut, stdErr, args) eg := &errgroup.Group{} eg.Go(func() error { defer cmd.SetIn(nil) return cmd.ExecuteContext(ctx) }) return eg } // Exec runs the provided cobra command with the given reader as STD_IN and the given args. // Returns STD_OUT, STD_ERR and the error from the execution. func Exec(t testing.TB, cmd *cobra.Command, stdIn io.Reader, args ...string) (string, string, error) { return ExecCtx(t.Context(), cmd, stdIn, args...) } func ExecCtx(ctx context.Context, cmd *cobra.Command, stdIn io.Reader, args ...string) (string, string, error) { stdOut, stdErr := &bytes.Buffer{}, &bytes.Buffer{} prepareCmd(cmd, stdIn, stdOut, stdErr, args) // needs to be on a separate line to ensure that the output buffers are read AFTER the command ran err := cmd.ExecuteContext(ctx) return stdOut.String(), stdErr.String(), err } // ExecNoErr is a helper that assumes a successful run from Exec. // Returns STD_OUT. func ExecNoErr(t testing.TB, cmd *cobra.Command, args ...string) string { return ExecNoErrCtx(t.Context(), t, cmd, args...) } func ExecNoErrCtx(ctx context.Context, t require.TestingT, cmd *cobra.Command, args ...string) string { stdOut, stdErr, err := ExecCtx(ctx, cmd, nil, args...) if err == nil { require.Len(t, stdErr, 0, "std_out: %s\nstd_err: %s", stdOut, stdErr) } else { require.ErrorIsf(t, err, context.Canceled, "std_out: %s\nstd_err: %s", stdOut, stdErr) } return stdOut } // ExecExpectedErr is a helper that assumes a failing run from Exec returning ErrNoPrintButFail // Returns STD_ERR. func ExecExpectedErr(t testing.TB, cmd *cobra.Command, args ...string) string { return ExecExpectedErrCtx(t.Context(), t, cmd, args...) } func ExecExpectedErrCtx(ctx context.Context, t require.TestingT, cmd *cobra.Command, args ...string) string { stdOut, stdErr, err := ExecCtx(ctx, cmd, nil, args...) require.True(t, errors.Is(err, ErrNoPrintButFail), "std_out: %s\nstd_err: %s", stdOut, stdErr) require.Len(t, stdOut, 0, stdErr) return stdErr } type CommandExecuter struct { New func() *cobra.Command Ctx context.Context PersistentArgs []string } func (c *CommandExecuter) Exec(stdin io.Reader, args ...string) (string, string, error) { return ExecCtx(c.Ctx, c.New(), stdin, append(c.PersistentArgs, args...)...) } func (c *CommandExecuter) ExecBackground(stdin io.Reader, stdOut, stdErr io.Writer, args ...string) *errgroup.Group { return ExecBackgroundCtx(c.Ctx, c.New(), stdin, stdOut, stdErr, append(c.PersistentArgs, args...)...) } func (c *CommandExecuter) ExecNoErr(t require.TestingT, args ...string) string { return ExecNoErrCtx(c.Ctx, t, c.New(), append(c.PersistentArgs, args...)...) } func (c *CommandExecuter) ExecExpectedErr(t require.TestingT, args ...string) string { return ExecExpectedErrCtx(c.Ctx, t, c.New(), append(c.PersistentArgs, args...)...) } type URL struct { url.URL } var _ pflag.Value = (*URL)(nil) func (u *URL) Set(s string) error { uu, err := url.Parse(s) if err != nil { return err } u.URL = *uu return nil } func (*URL) Type() string { return "url" }