chore: cleanup tests

Signed-off-by: Alano Terblanche <18033717+Benehiko@users.noreply.github.com>

cli/command/builder/prune_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"testing"
-	"time"
 
 	"github.com/docker/cli/internal/test"
 	"github.com/docker/docker/api/types"
@@ -19,8 +18,6 @@ func TestBuilderPromptTermination(t *testing.T) {
 			return nil, errors.New("fakeClient builderPruneFunc should not be called")
 		},
 	})
-	// set a fake reader so that our kill signal reaches the prompt before the prompt reads from stdin
-	cli.SetIn(test.NewFakeStreamIn(ctx, time.Second*2))
 	cmd := NewPruneCommand(cli)
 	test.TerminatePrompt(ctx, t, cmd, cli, nil)
 }

cli/command/container/prune_test.go

@@ -3,7 +3,6 @@ package container
 import (
 	"context"
 	"testing"
-	"time"
 
 	"github.com/docker/cli/internal/test"
 	"github.com/docker/docker/api/types"
@@ -20,9 +19,6 @@ func TestContainerPrunePromptTermination(t *testing.T) {
 			return types.ContainersPruneReport{}, errors.New("fakeClient containerPruneFunc should not be called")
 		},
 	})
-
-	// set a fake reader so that our kill signal reaches the prompt before the prompt reads from stdin
-	cli.SetIn(test.NewFakeStreamIn(ctx, time.Second*2))
 	cmd := NewPruneCommand(cli)
 	test.TerminatePrompt(ctx, t, cmd, cli, nil)
 }

cli/command/image/prune_test.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"io"
 	"testing"
-	"time"
 
 	"github.com/docker/cli/internal/test"
 	"github.com/docker/docker/api/types"
@@ -113,8 +112,6 @@ func TestPrunePromptTermination(t *testing.T) {
 			return types.ImagesPruneReport{}, errors.New("fakeClient imagesPruneFunc should not be called")
 		},
 	})
-	// set a fake reader so that our kill signal reaches the prompt before the prompt reads from stdin
-	cli.SetIn(test.NewFakeStreamIn(ctx, time.Second*2))
 	cmd := NewPruneCommand(cli)
 	test.TerminatePrompt(ctx, t, cmd, cli, nil)
 }

cli/command/network/prune_test.go

@@ -3,7 +3,6 @@ package network
 import (
 	"context"
 	"testing"
-	"time"
 
 	"github.com/docker/cli/internal/test"
 	"github.com/docker/docker/api/types"
@@ -20,8 +19,6 @@ func TestNetworkPrunePromptTermination(t *testing.T) {
 			return types.NetworksPruneReport{}, errors.New("fakeClient networkPruneFunc should not be called")
 		},
 	})
-	// set a fake reader so that our kill signal reaches the prompt before the prompt reads from stdin
-	cli.SetIn(test.NewFakeStreamIn(ctx, time.Second*2))
 	cmd := NewPruneCommand(cli)
 	test.TerminatePrompt(ctx, t, cmd, cli, nil)
 }

cli/command/network/remove_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"io"
 	"testing"
-	"time"
 
 	"github.com/docker/cli/internal/test"
 	"github.com/docker/docker/api/types"
@@ -113,8 +112,6 @@ func TestNetworkRemovePromptTermination(t *testing.T) {
 			}, nil, nil
 		},
 	})
-	// set a fake reader so that our kill signal reaches the prompt before the prompt reads from stdin
-	cli.SetIn(test.NewFakeStreamIn(ctx, time.Second*2))
 	cmd := newRemoveCommand(cli)
 	cmd.SetArgs([]string{"existing-network"})
 	test.TerminatePrompt(ctx, t, cmd, cli, nil)

cli/command/plugin/upgrade_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"io"
 	"testing"
-	"time"
 
 	"github.com/docker/cli/internal/test"
 	"github.com/docker/docker/api/types"
@@ -14,9 +13,6 @@ import (
 )
 
 func TestUpgradePromptTermination(t *testing.T) {
-	t.Cleanup(func() {
-	})
-
 	ctx, cancel := context.WithCancel(context.Background())
 	t.Cleanup(cancel)
 
@@ -33,7 +29,6 @@ func TestUpgradePromptTermination(t *testing.T) {
 			}, []byte{}, nil
 		},
 	})
-	cli.SetIn(test.NewFakeStreamIn(ctx, time.Second*2))
 	cmd := newUpgradeCommand(cli)
 	// need to set a remote address that does not match the plugin
 	// reference sent by the `pluginInspectFunc`

cli/command/system/prune_test.go

@@ -3,7 +3,6 @@ package system
 import (
 	"context"
 	"testing"
-	"time"
 
 	"github.com/docker/cli/cli/config/configfile"
 	"github.com/docker/cli/internal/test"
@@ -68,8 +67,6 @@ func TestSystemPrunePromptTermination(t *testing.T) {
 		},
 	})
 
-	// set a fake reader so that our kill signal reaches the prompt before the prompt reads from stdin
-	cli.SetIn(test.NewFakeStreamIn(ctx, time.Second*2))
 	cmd := newPruneCommand(cli)
 	test.TerminatePrompt(ctx, t, cmd, cli, nil)
 }

cli/command/trust/revoke_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"io"
 	"testing"
-	"time"
 
 	"github.com/docker/cli/cli/trust"
 	"github.com/docker/cli/internal/test"
@@ -157,10 +156,8 @@ func TestRevokeTrustPromptTermination(t *testing.T) {
 	t.Cleanup(cancel)
 
 	cli := test.NewFakeCli(&fakeClient{})
-	cli.SetIn(test.NewFakeStreamIn(ctx, 2*time.Second))
 	cmd := newRevokeCommand(cli)
 	cmd.SetArgs([]string{"example/trust-demo"})
-
 	test.TerminatePrompt(ctx, t, cmd, cli, nil)
 	assert.Equal(t, cli.ErrBuffer().String(), "")
 	golden.Assert(t, cli.OutBuffer().String(), "trust-revoke-prompt-termination.golden")

cli/command/utils_test.go

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/docker/cli/cli/command"
+	"github.com/docker/cli/cli/streams"
 	"github.com/docker/cli/internal/test"
 	"github.com/pkg/errors"
 	"gotest.tools/v3/assert"
@@ -81,46 +82,98 @@ func TestPromptForConfirmation(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	t.Cleanup(cancel)
 
+	t.Run("case=prompt should wait for user input", func(t *testing.T) {
+		buf := new(bytes.Buffer)
+		wroteHook := make(chan struct{}, 1)
+		bufioWriter := bufio.NewWriter(buf)
+		w := test.NewWriterWithHook(bufioWriter, func(p []byte) {
+			wroteHook <- struct{}{}
+		})
+
+		r, _, err := os.Pipe()
+		assert.NilError(t, err)
+		in := streams.NewIn(r)
+
+		result := make(chan bool, 1)
+		go func() {
+			result <- command.PromptForConfirmation(ctx, in, w, "")
+		}()
+
+		// wait for the prompt to write to the buffer
+		pollForPromptOutput(ctx, t, wroteHook)
+
+		assert.NilError(t, bufioWriter.Flush())
+		assert.Check(t, strings.Contains(buf.String(), "Are you sure you want to proceed? [y/N]"))
+
+		writeCtx, writeCancel := context.WithTimeout(ctx, 100*time.Millisecond)
+		defer writeCancel()
+
+		writeInputChan := make(chan struct{})
+
+		go func() {
+			_, err := w.Write([]byte("n"))
+			assert.NilError(t, err)
+			writeInputChan <- struct{}{}
+		}()
+
+		select {
+		case <-writeCtx.Done():
+			t.Fatal("PromptForConfirmation did not receive user input")
+		case <-result:
+			t.Fatal("PromptForConfirmation returned before user input")
+		case <-writeInputChan:
+		}
+	})
+
 	t.Run("case=terminate the prompt with SIGINT", func(t *testing.T) {
 		wroteHook := make(chan struct{}, 1)
 		buf := new(bytes.Buffer)
 		bufioWriter := bufio.NewWriter(buf)
-		w := test.NewWriterWithHook(bufioWriter, func() {
+		w := test.NewWriterWithHook(bufioWriter, func(p []byte) {
 			wroteHook <- struct{}{}
 		})
 
-		in := test.NewFakeStreamIn(ctx, time.Second*2)
-		t.Cleanup(func() {
-			assert.NilError(t, in.Close())
-		})
+		r, _, err := os.Pipe()
+		assert.NilError(t, err)
+		in := streams.NewIn(r)
+
 		result := make(chan bool, 1)
+		defer close(result)
+
 		go func() {
-			defer close(result)
 			result <- command.PromptForConfirmation(ctx, in, w, "")
 		}()
 
 		// wait for the Prompt to write to the buffer
 		pollForPromptOutput(ctx, t, wroteHook)
+
 		assert.NilError(t, bufioWriter.Flush())
 		assert.Equal(t, strings.TrimSpace(buf.String()), "Are you sure you want to proceed? [y/N]")
 
+		resultCtx, resultCancel := context.WithTimeout(ctx, 100*time.Millisecond)
+		defer resultCancel()
+
 		syscall.Kill(syscall.Getpid(), syscall.SIGINT)
+
 		select {
+		case <-resultCtx.Done():
+			t.Fatal("PromptForConfirmation did not return after SIGINT")
 		case r := <-result:
 			assert.Check(t, !r)
-		case <-time.After(100 * time.Millisecond):
-			t.Fatal("PromptForConfirmation did not return after SIGINT")
 		}
 	})
 
 	t.Run("case=prompt should return on io.EOF", func(t *testing.T) {
 		buf := new(bytes.Buffer)
 		wroteHook := make(chan struct{}, 1)
 		bufioWriter := bufio.NewWriter(buf)
-		w := test.NewWriterWithHook(bufioWriter, func() {
+		w := test.NewWriterWithHook(bufioWriter, func(p []byte) {
 			wroteHook <- struct{}{}
 		})
-		in := test.NewFakeStreamIn(ctx, 0)
+
+		r, _, err := os.Pipe()
+		assert.NilError(t, err)
+		in := streams.NewIn(r)
 
 		result := make(chan bool, 1)
 		go func() {
@@ -132,13 +185,17 @@ func TestPromptForConfirmation(t *testing.T) {
 
 		assert.NilError(t, bufioWriter.Flush())
 		assert.Check(t, strings.Contains(buf.String(), "Are you sure you want to proceed? [y/N]"))
+
 		assert.NilError(t, in.Close())
 
+		resultCtx, resultCancel := context.WithTimeout(ctx, 100*time.Millisecond)
+		defer resultCancel()
+
 		select {
+		case <-resultCtx.Done():
+			t.Fatal("PromptForConfirmation did not return after io.EOF")
 		case r := <-result:
 			assert.Check(t, !r)
-		case <-time.After(100 * time.Millisecond):
-			t.Fatal("PromptForConfirmation did not return after io.EOF")
 		}
 	})
 }

cli/command/volume/prune_test.go

@@ -7,7 +7,6 @@ import (
 	"runtime"
 	"strings"
 	"testing"
-	"time"
 
 	"github.com/docker/cli/cli/streams"
 	"github.com/docker/cli/internal/test"
@@ -196,8 +195,6 @@ func TestVolumePrunePromptTerminate(t *testing.T) {
 		},
 	})
 
-	// set a fake reader so that our kill signal reaches the prompt before the prompt reads from stdin
-	cli.SetIn(test.NewFakeStreamIn(ctx, time.Second*2))
 	cmd := NewPruneCommand(cli)
 	test.TerminatePrompt(ctx, t, cmd, cli, nil)
 

internal/test/cmd.go

@@ -2,10 +2,12 @@ package test
 
 import (
 	"context"
+	"os"
 	"syscall"
 	"testing"
 	"time"
 
+	"github.com/docker/cli/cli/streams"
 	"github.com/spf13/cobra"
 	"gotest.tools/v3/assert"
 )
@@ -16,24 +18,57 @@ func TerminatePrompt(ctx context.Context, t *testing.T, cmd *cobra.Command, cli
 	errChan := make(chan error)
 	defer close(errChan)
 
+	// wrap the out stream to detect when the prompt is ready
+	writerHookChan := make(chan struct{})
+	defer close(writerHookChan)
+
+	outStream := streams.NewOut(NewWriterWithHook(cli.OutBuffer(), func(p []byte) {
+		writerHookChan <- struct{}{}
+	}))
+	cli.SetOut(outStream)
+
+	r, _, err := os.Pipe()
+	assert.NilError(t, err)
+	cli.SetIn(streams.NewIn(r))
+
 	go func() {
 		errChan <- cmd.ExecuteContext(ctx)
 	}()
 
-	time.Sleep(100 * time.Millisecond)
+	writeCtx, writeCancel := context.WithTimeout(ctx, 100*time.Millisecond)
+	defer writeCancel()
+
+	// wait for the prompt to be ready
+	select {
+	case <-writeCtx.Done():
+		t.Fatalf("command %s did not write prompt to stdout", cmd.Name())
+	case <-writerHookChan:
+		// drain the channel for future buffer writes
+		go func() {
+			for range writerHookChan {
+			}
+		}()
+	}
+
+	assert.Check(t, cli.OutBuffer().Len() > 0)
 
 	syscall.Kill(syscall.Getpid(), syscall.SIGINT)
 	// sigint and sigterm are caught by the prompt
 	// this allows us to gracefully exit the prompt with a 0 exit code
+	errCtx, errCancel := context.WithTimeout(ctx, 100*time.Millisecond)
+	defer errCancel()
+
 	select {
+	case <-errCtx.Done():
+		t.Logf("command stdout:\n%s\n", cli.OutBuffer().String())
+		t.Logf("command stderr:\n%s\n", cli.ErrBuffer().String())
+		t.Fatalf("command %s did not return after SIGINT", cmd.Name())
 	case err := <-errChan:
 		if assertFunc != nil {
 			assertFunc(t, err)
 			return
 		}
 		assert.NilError(t, err)
-	case <-time.After(1000 * time.Millisecond):
-		t.Fatalf("command %s did not return after SIGINT", cmd.Name())
 	}
 
 	assert.Equal(t, cli.ErrBuffer().String(), "")