From 4a17a04514e754aded1594d4125264e1e1976c6b Mon Sep 17 00:00:00 2001
From: antichris <881420+antichris@users.noreply.github.com>
Date: Tue, 6 Apr 2021 06:46:20 +0300
Subject: [PATCH] complete: enable complete.Complete() output capturing (#138)

* complete: test capturing Complete() output

Implement an Example test to demonstrate capturing the output
of Complete(), which is crucial for integration tests.

* complete: do not hard-code the I/O streams at the package initialization

Instead of defining the input/output streams as unexported global vars
that only get their values assigned to once, at the very initialization
of the package, use the values that `os.Stdin` and `os.Stdout` have at
the particular moment on every `complete.Complete()` call.

Fix #137

* complete: capture and discard output in TestComplete

Restore earlier behavior using proper stream redirection this time.

* complete: output capturing example: define things in the package scope

Define the `stringLookup` func type and `promptEnv` func in the package
scope instead of the `ExampleComplete_outputCapturing` test.

* complete: rename the `stringLookup` func type to `getEnvFn`
---
 complete.go      | 10 ++++++----
 complete_test.go | 48 ++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/complete.go b/complete.go
index 2a4c905..0641c39 100644
--- a/complete.go
+++ b/complete.go
@@ -49,10 +49,8 @@ func (p PredictFunc) Predict(prefix string) []string {
 }
 
 var (
-	getEnv           = os.Getenv
-	exit             = os.Exit
-	out    io.Writer = os.Stdout
-	in     io.Reader = os.Stdin
+	getEnv = os.Getenv
+	exit   = os.Exit
 )
 
 // Complete the command line arguments for the given command in the case that the program
@@ -66,6 +64,10 @@ func Complete(name string, cmd Completer) {
 		doUninstall = getEnv("COMP_UNINSTALL") == "1"
 		yes         = getEnv("COMP_YES") == "1"
 	)
+	var (
+		out io.Writer = os.Stdout
+		in  io.Reader = os.Stdin
+	)
 	if doInstall || doUninstall {
 		install.Run(name, doUninstall, yes, out, in)
 		exit(0)
diff --git a/complete_test.go b/complete_test.go
index 3a0809f..9870910 100644
--- a/complete_test.go
+++ b/complete_test.go
@@ -1,8 +1,10 @@
 package complete
 
 import (
+	"io"
 	"io/ioutil"
 	"os"
+	"strconv"
 	"testing"
 
 	"github.com/posener/complete/v2/internal/arg"
@@ -159,9 +161,17 @@ func TestComplete(t *testing.T) {
 	defer func() {
 		getEnv = os.Getenv
 		exit = os.Exit
-		out = os.Stdout
 	}()
 
+	in, out, err := os.Pipe()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func(o *os.File) { os.Stdout = o }(os.Stdout)
+	defer out.Close()
+	os.Stdout = out
+	go io.Copy(ioutil.Discard, in)
+
 	tests := []struct {
 		line, point string
 		shouldExit  bool
@@ -203,7 +213,6 @@ func TestComplete(t *testing.T) {
 			exit = func(int) {
 				isExit = true
 			}
-			out = ioutil.Discard
 			if tt.shouldPanic {
 				assert.Panics(t, func() { testCmd.Complete("") })
 			} else {
@@ -214,6 +223,24 @@ func TestComplete(t *testing.T) {
 	}
 }
 
+// ExampleComplete_outputCapturing demonstrates the ability to capture
+// the output of Complete() invocations, crucial for integration tests.
+func ExampleComplete_outputCapturing() {
+	defer func(f func(int)) { exit = f }(exit)
+	defer func(f getEnvFn) { getEnv = f }(getEnv)
+	exit = func(int) {}
+
+	// This is where the actual example starts:
+
+	cmd := &Command{Sub: map[string]*Command{"bar": {}}}
+	getEnv = promptEnv("foo b")
+
+	Complete("foo", cmd)
+
+	// Output:
+	// bar
+}
+
 type set []string
 
 func (s set) Predict(_ string) []string {
@@ -245,3 +272,20 @@ func TestHasPrefix(t *testing.T) {
 		})
 	}
 }
+
+// getEnvFn emulates os.GetEnv by mapping one string to another.
+type getEnvFn = func(string) string
+
+// promptEnv returns getEnvFn that emulates the environment variables
+// a shell would set when its prompt has the given contents.
+var promptEnv = func(contents string) getEnvFn {
+	return func(key string) string {
+		switch key {
+		case "COMP_LINE":
+			return contents
+		case "COMP_POINT":
+			return strconv.Itoa(len(contents))
+		}
+		return ""
+	}
+}