feat: allow setting stdin, stdout and stderr in retry command (#676)

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 - Migrate from github.com/jetstack/cert-manager to github.com/cert-manager/cert-manager
   [#669](https://github.com/Kong/kubernetes-testing-framework/pull/669)
+- Allow setting stdin, stdout, stderr in `retry.Command`
+  [#676](https://github.com/Kong/kubernetes-testing-framework/pull/676)
 
 ## v0.31.0
 

internal/retry/command.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"os/exec"
 	"time"
 
@@ -23,24 +24,42 @@ type Doer interface {
 
 type ErrorFunc func(error, *bytes.Buffer, *bytes.Buffer) error
 
-type commandDoer struct {
-	cmd  string
-	args []string
+type CommandDoer struct {
+	stdin  io.Reader
+	stdout io.Writer
+	stderr io.Writer
+	cmd    string
+	args   []string
 }
 
-func Command(cmd string, args ...string) Doer {
-	return commandDoer{
+func Command(cmd string, args ...string) CommandDoer {
+	return CommandDoer{
 		cmd:  cmd,
 		args: args,
 	}
 }
 
-func (c commandDoer) Do(ctx context.Context) error {
+func (c CommandDoer) WithStdin(r io.Reader) CommandDoer {
+	c.stdin = r
+	return c
+}
+
+func (c CommandDoer) WithStdout(w io.Writer) CommandDoer {
+	c.stdout = w
+	return c
+}
+
+func (c CommandDoer) WithStderr(w io.Writer) CommandDoer {
+	c.stderr = w
+	return c
+}
+
+func (c CommandDoer) Do(ctx context.Context) error {
 	return retry.Do(func() error {
 		cmd, stdout, stderr := c.createCmd(ctx)
 		err := cmd.Run()
 		if err != nil {
-			return fmt.Errorf("command %q with args [%v] failed STDOUT=(%s) STDERR=(%s): %w",
+			return fmt.Errorf("command %q with args %v failed STDOUT=(%s) STDERR=(%s): %w",
 				c.cmd, c.args, stdout.String(), stderr.String(), err,
 			)
 		}
@@ -52,7 +71,7 @@ func (c commandDoer) Do(ctx context.Context) error {
 
 // DoWithErrorHandling executes the command and runs errorFunc passing a resulting err, stdout and stderr to be handled
 // by the caller. The errorFunc is going to be called only when the resulting err != nil.
-func (c commandDoer) DoWithErrorHandling(ctx context.Context, errorFunc ErrorFunc) error {
+func (c CommandDoer) DoWithErrorHandling(ctx context.Context, errorFunc ErrorFunc) error {
 	return retry.Do(func() error {
 		cmd, stdout, stderr := c.createCmd(ctx)
 		err := cmd.Run()
@@ -65,16 +84,30 @@ func (c commandDoer) DoWithErrorHandling(ctx context.Context, errorFunc ErrorFun
 	)
 }
 
-func (c commandDoer) createCmd(ctx context.Context) (*exec.Cmd, *bytes.Buffer, *bytes.Buffer) {
+func (c CommandDoer) createCmd(ctx context.Context) (*exec.Cmd, *bytes.Buffer, *bytes.Buffer) {
 	stdout := new(bytes.Buffer)
+	if c.stdout == nil {
+		c.stdout = stdout
+	} else {
+		c.stdout = io.MultiWriter(c.stdout, stdout)
+	}
+
 	stderr := new(bytes.Buffer)
+	if c.stderr == nil {
+		c.stderr = stderr
+	} else {
+		c.stderr = io.MultiWriter(c.stderr, stderr)
+	}
+
 	cmd := exec.CommandContext(ctx, c.cmd, c.args...) //nolint:gosec
-	cmd.Stdout = stdout
-	cmd.Stderr = stderr
+	cmd.Stdin = c.stdin
+	cmd.Stdout = c.stdout
+	cmd.Stderr = c.stderr
+
 	return cmd, stdout, stderr
 }
 
-func (c commandDoer) createOpts(ctx context.Context) []retry.Option {
+func (c CommandDoer) createOpts(ctx context.Context) []retry.Option {
 	return []retry.Option{
 		retry.Context(ctx),
 		retry.Delay(retryWait),

internal/retry/command_test.go

@@ -40,3 +40,33 @@ func TestDoWithErrorHandling(t *testing.T) {
 		require.True(t, wasCalled, "expected errorFunc to be called because the command has failed")
 	})
 }
+
+func TestDo(t *testing.T) {
+	t.Run("passing stdout works", func(t *testing.T) {
+		stdout := &bytes.Buffer{}
+		cmd := retry.Command("echo", "hello").WithStdout(stdout)
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+
+		err := cmd.Do(ctx)
+		require.NoError(t, err)
+		require.Equal(t, "hello\n", stdout.String())
+	})
+
+	t.Run("passing stdin works", func(t *testing.T) {
+		stdin := bytes.NewBufferString("hello")
+		cmd := retry.Command("cat").WithStdin(stdin)
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+
+		err := cmd.Do(ctx)
+		require.NoError(t, err)
+	})
+
+	// Testing stderr might not be reliable because it's not guaranteed that
+	// the command will fail in the time we allow it to run.
+	// Alternative would be to wait long enough so that we're it ran but that
+	// would unnecessarily make the tests longer.
+}

pkg/clusters/addons/metallb/metallb.go

@@ -1,15 +1,13 @@
 package metallb
 
 import (
-	"bytes"
 	"context"
 	"crypto/rand"
 	"encoding/base64"
 	"fmt"
 	"io"
 	"net"
 	"os"
-	"os/exec"
 	"time"
 
 	corev1 "k8s.io/api/core/v1"
@@ -23,6 +21,7 @@ import (
 	kustomize "sigs.k8s.io/kustomize/api/types"
 	"sigs.k8s.io/kustomize/kyaml/resid"
 
+	"github.com/kong/kubernetes-testing-framework/internal/retry"
 	"github.com/kong/kubernetes-testing-framework/pkg/clusters"
 	"github.com/kong/kubernetes-testing-framework/pkg/clusters/types/kind"
 	"github.com/kong/kubernetes-testing-framework/pkg/utils/docker"
@@ -112,7 +111,7 @@ func (a *addon) Delete(ctx context.Context, cluster clusters.Cluster) error {
 	}
 	defer os.Remove(kubeconfig.Name())
 
-	return metallbDeleteHack(kubeconfig)
+	return metallbDeleteHack(ctx, kubeconfig)
 }
 
 func (a *addon) Ready(ctx context.Context, cluster clusters.Cluster) ([]runtime.Object, bool, error) {
@@ -164,7 +163,7 @@ func deployMetallbForKindCluster(ctx context.Context, cluster clusters.Cluster,
 
 	// create the metallb deployment and related resources (do this first so that
 	// we can create the IPAddressPool below with its CRD already in place).
-	if err := metallbDeployHack(cluster); err != nil {
+	if err := metallbDeployHack(ctx, cluster); err != nil {
 		return fmt.Errorf("failed to deploy metallb: %w", err)
 	}
 
@@ -346,7 +345,7 @@ func getManifest() (io.Reader, error) {
 	})
 }
 
-func metallbDeployHack(cluster clusters.Cluster) error {
+func metallbDeployHack(ctx context.Context, cluster clusters.Cluster) error {
 	// generate a temporary kubeconfig since we're going to be using kubectl
 	kubeconfig, err := clusters.TempKubeconfig(cluster)
 	if err != nil {
@@ -365,20 +364,14 @@ func metallbDeployHack(cluster clusters.Cluster) error {
 		return fmt.Errorf("could not deploy metallb: %w", err)
 	}
 
-	stderr := new(bytes.Buffer)
-	cmd := exec.Command("kubectl", deployArgs...)
-	cmd.Stdout = io.Discard
-	cmd.Stderr = stderr
-	cmd.Stdin = manifest
-
-	if err := cmd.Run(); err != nil {
-		return fmt.Errorf("%s: %w", stderr.String(), err)
-	}
-
-	return nil
+	// ensure the repo exists
+	return retry.Command("kubectl", deployArgs...).
+		WithStdin(manifest).
+		WithStdout(io.Discard).
+		Do(ctx)
 }
 
-func metallbDeleteHack(kubeconfig *os.File) error {
+func metallbDeleteHack(ctx context.Context, kubeconfig *os.File) error {
 	deployArgs := []string{
 		"--kubeconfig", kubeconfig.Name(),
 		"delete", "-f", "-",
@@ -389,15 +382,8 @@ func metallbDeleteHack(kubeconfig *os.File) error {
 		return fmt.Errorf("could not delete metallb: %w", err)
 	}
 
-	stderr := new(bytes.Buffer)
-	cmd := exec.Command("kubectl", deployArgs...)
-	cmd.Stdout = io.Discard
-	cmd.Stderr = stderr
-	cmd.Stdin = manifest
-
-	if err := cmd.Run(); err != nil {
-		return fmt.Errorf("%s: %w", stderr.String(), err)
-	}
-
-	return nil
+	return retry.Command("kubectl", deployArgs...).
+		WithStdin(manifest).
+		WithStdout(io.Discard).
+		Do(ctx)
 }