General refactoring & cleanup

Signed-off-by: Dave Henderson <dhenderson@gmail.com>

1 files changed, 17 insertions, 3 deletions

diff --git a/data/datasource_stdin.go b/data/datasource_stdin.go
index 007d935b..b8bd0eff 100644
--- a/data/datasource_stdin.go
+++ b/data/datasource_stdin.go
@@ -4,17 +4,31 @@ import (
 	"context"
 	"io"
 	"io/ioutil"
+	"os"
 
 	"github.com/pkg/errors"
 )
 
-// stdin - for overriding in tests
-var stdin io.Reader
-
 func readStdin(ctx context.Context, source *Source, args ...string) ([]byte, error) {
+	stdin := stdinFromContext(ctx)
+
 	b, err := ioutil.ReadAll(stdin)
 	if err != nil {
 		return nil, errors.Wrapf(err, "Can't read %s", stdin)
 	}
 	return b, nil
 }
+
+type stdinCtxKey struct{}
+
+func ContextWithStdin(ctx context.Context, r io.Reader) context.Context {
+	return context.WithValue(ctx, stdinCtxKey{}, r)
+}
+
+func stdinFromContext(ctx context.Context) io.Reader {
+	if r, ok := ctx.Value(stdinCtxKey{}).(io.Reader); ok {
+		return r
+	}
+
+	return os.Stdin
+}