General refactoring & cleanup

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

6 files changed, 55 insertions, 39 deletions

diff --git a/data/datasource.go b/data/datasource.go
index 539ba122..c010f89d 100644
--- a/data/datasource.go
+++ b/data/datasource.go
@@ -78,8 +78,9 @@ func (d *Data) lookupReader(scheme string) (func(context.Context, *Source, ...st
 }
 
 // Data -
+// Deprecated: will be replaced in future
 type Data struct {
-	ctx context.Context
+	Ctx context.Context
 
 	Sources map[string]*Source
 
@@ -87,7 +88,7 @@ type Data struct {
 	cache         map[string][]byte
 
 	// headers from the --datasource-header/-H option that don't reference datasources from the commandline
-	extraHeaders map[string]http.Header
+	ExtraHeaders map[string]http.Header
 }
 
 // Cleanup - clean up datasources before shutting the process down - things
@@ -114,41 +115,41 @@ func NewData(datasourceArgs, headerArgs []string) (*Data, error) {
 func FromConfig(ctx context.Context, cfg *config.Config) *Data {
 	// XXX: This is temporary, and will be replaced with something a bit cleaner
 	// when datasources are refactored
-	stdin = cfg.Stdin
+	ctx = ContextWithStdin(ctx, cfg.Stdin)
 
 	sources := map[string]*Source{}
 	for alias, d := range cfg.DataSources {
 		sources[alias] = &Source{
 			Alias:  alias,
 			URL:    d.URL,
-			header: d.Header,
+			Header: d.Header,
 		}
 	}
 	for alias, d := range cfg.Context {
 		sources[alias] = &Source{
 			Alias:  alias,
 			URL:    d.URL,
-			header: d.Header,
+			Header: d.Header,
 		}
 	}
 	return &Data{
-		ctx:          ctx,
+		Ctx:          ctx,
 		Sources:      sources,
-		extraHeaders: cfg.ExtraHeaders,
+		ExtraHeaders: cfg.ExtraHeaders,
 	}
 }
 
 // Source - a data source
 type Source struct {
+	Alias             string
 	URL               *url.URL
+	Header            http.Header             // used for http[s]: URLs, nil otherwise
 	fs                afero.Fs                // used for file: URLs, nil otherwise
 	hc                *http.Client            // used for http[s]: URLs, nil otherwise
 	vc                *vault.Vault            // used for vault: URLs, nil otherwise
 	kv                *libkv.LibKV            // used for consul:, etcd:, zookeeper: & boltdb: URLs, nil otherwise
 	asmpg             awssmpGetter            // used for aws+smp:, nil otherwise
 	awsSecretsManager awsSecretsManagerGetter // used for aws+sm, nil otherwise
-	header            http.Header             // used for http[s]: URLs, nil otherwise
-	Alias             string
 	mediaType         string
 }
 
@@ -246,7 +247,7 @@ func (d *Data) DefineDatasource(alias, value string) (string, error) {
 	s := &Source{
 		Alias:  alias,
 		URL:    srcURL,
-		header: d.extraHeaders[alias],
+		Header: d.ExtraHeaders[alias],
 	}
 	if d.Sources == nil {
 		d.Sources = make(map[string]*Source)
@@ -271,7 +272,7 @@ func (d *Data) lookupSource(alias string) (*Source, error) {
 		source = &Source{
 			Alias:  alias,
 			URL:    srcURL,
-			header: d.extraHeaders[alias],
+			Header: d.ExtraHeaders[alias],
 		}
 		d.Sources[alias] = source
 	}
@@ -304,13 +305,13 @@ func (d *Data) readDataSource(ctx context.Context, alias string, args ...string)
 
 // Include -
 func (d *Data) Include(alias string, args ...string) (string, error) {
-	data, _, err := d.readDataSource(d.ctx, alias, args...)
+	data, _, err := d.readDataSource(d.Ctx, alias, args...)
 	return data, err
 }
 
 // Datasource -
 func (d *Data) Datasource(alias string, args ...string) (interface{}, error) {
-	data, mimeType, err := d.readDataSource(d.ctx, alias, args...)
+	data, mimeType, err := d.readDataSource(d.Ctx, alias, args...)
 	if err != nil {
 		return nil, err
 	}
@@ -355,7 +356,7 @@ func (d *Data) DatasourceReachable(alias string, args ...string) bool {
 	if !ok {
 		return false
 	}
-	_, err := d.readSource(d.ctx, source, args...)
+	_, err := d.readSource(d.Ctx, source, args...)
 	return err == nil
 }
 
diff --git a/data/datasource_http.go b/data/datasource_http.go
index ed1c97b7..03e1ac4c 100644
--- a/data/datasource_http.go
+++ b/data/datasource_http.go
@@ -34,7 +34,7 @@ func readHTTP(ctx context.Context, source *Source, args ...string) ([]byte, erro
 	if err != nil {
 		return nil, err
 	}
-	req.Header = source.header
+	req.Header = source.Header
 	res, err := source.hc.Do(req)
 	if err != nil {
 		return nil, err
diff --git a/data/datasource_http_test.go b/data/datasource_http_test.go
index e2c13a8f..5d279712 100644
--- a/data/datasource_http_test.go
+++ b/data/datasource_http_test.go
@@ -58,7 +58,7 @@ func TestHTTPFile(t *testing.T) {
 		hc: client,
 	}
 	data := &Data{
-		ctx:     context.Background(),
+		Ctx:     context.Background(),
 		Sources: sources,
 	}
 
@@ -88,7 +88,7 @@ func TestHTTPFileWithHeaders(t *testing.T) {
 			Path:   "/foo",
 		},
 		hc: client,
-		header: http.Header{
+		Header: http.Header{
 			"Foo":             {"bar"},
 			"foo":             {"baz"},
 			"User-Agent":      {},
@@ -96,7 +96,7 @@ func TestHTTPFileWithHeaders(t *testing.T) {
 		},
 	}
 	data := &Data{
-		ctx:     context.Background(),
+		Ctx:     context.Background(),
 		Sources: sources,
 	}
 	expected := http.Header{
@@ -113,9 +113,9 @@ func TestHTTPFileWithHeaders(t *testing.T) {
 		"User-Agent":      {"Go-http-client/1.1"},
 	}
 	data = &Data{
-		ctx:          context.Background(),
+		Ctx:          context.Background(),
 		Sources:      sources,
-		extraHeaders: map[string]http.Header{server.URL: expected},
+		ExtraHeaders: map[string]http.Header{server.URL: expected},
 	}
 	actual, err = data.Datasource(server.URL)
 	assert.NoError(t, err)
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
+}
diff --git a/data/datasource_stdin_test.go b/data/datasource_stdin_test.go
index 1b92914e..6cd34133 100644
--- a/data/datasource_stdin_test.go
+++ b/data/datasource_stdin_test.go
@@ -11,15 +11,12 @@ import (
 func TestReadStdin(t *testing.T) {
 	ctx := context.Background()
 
-	defer func() {
-		stdin = nil
-	}()
-	stdin = strings.NewReader("foo")
+	ctx = ContextWithStdin(ctx, strings.NewReader("foo"))
 	out, err := readStdin(ctx, nil)
 	assert.NoError(t, err)
 	assert.Equal(t, []byte("foo"), out)
 
-	stdin = errorReader{}
+	ctx = ContextWithStdin(ctx, errorReader{})
 	_, err = readStdin(ctx, nil)
 	assert.Error(t, err)
 }
diff --git a/data/datasource_test.go b/data/datasource_test.go
index aadcef0e..5badec53 100644
--- a/data/datasource_test.go
+++ b/data/datasource_test.go
@@ -28,17 +28,17 @@ func TestNewData(t *testing.T) {
 	d, err = NewData([]string{"foo=http:///foo.json"}, []string{})
 	assert.NoError(t, err)
 	assert.Equal(t, "/foo.json", d.Sources["foo"].URL.Path)
-	assert.Empty(t, d.Sources["foo"].header)
+	assert.Empty(t, d.Sources["foo"].Header)
 
 	d, err = NewData([]string{"foo=http:///foo.json"}, []string{"bar=Accept: blah"})
 	assert.NoError(t, err)
 	assert.Equal(t, "/foo.json", d.Sources["foo"].URL.Path)
-	assert.Empty(t, d.Sources["foo"].header)
+	assert.Empty(t, d.Sources["foo"].Header)
 
 	d, err = NewData([]string{"foo=http:///foo.json"}, []string{"foo=Accept: blah"})
 	assert.NoError(t, err)
 	assert.Equal(t, "/foo.json", d.Sources["foo"].URL.Path)
-	assert.Equal(t, "blah", d.Sources["foo"].header["Accept"][0])
+	assert.Equal(t, "blah", d.Sources["foo"].Header["Accept"][0])
 }
 
 func TestDatasource(t *testing.T) {
@@ -344,12 +344,14 @@ func TestMimeTypeWithArg(t *testing.T) {
 
 func TestFromConfig(t *testing.T) {
 	ctx := context.Background()
+
 	cfg := &config.Config{}
+	actual := FromConfig(ctx, cfg)
 	expected := &Data{
-		ctx:     ctx,
+		Ctx:     actual.Ctx,
 		Sources: map[string]*Source{},
 	}
-	assert.EqualValues(t, expected, FromConfig(ctx, cfg))
+	assert.EqualValues(t, expected, actual)
 
 	cfg = &config.Config{
 		DataSources: map[string]config.DataSource{
@@ -358,8 +360,9 @@ func TestFromConfig(t *testing.T) {
 			},
 		},
 	}
+	actual = FromConfig(ctx, cfg)
 	expected = &Data{
-		ctx: ctx,
+		Ctx: actual.Ctx,
 		Sources: map[string]*Source{
 			"foo": {
 				Alias: "foo",
@@ -367,7 +370,7 @@ func TestFromConfig(t *testing.T) {
 			},
 		},
 	}
-	assert.EqualValues(t, expected, FromConfig(ctx, cfg))
+	assert.EqualValues(t, expected, actual)
 
 	cfg = &config.Config{
 		DataSources: map[string]config.DataSource{
@@ -389,8 +392,9 @@ func TestFromConfig(t *testing.T) {
 			},
 		},
 	}
+	actual = FromConfig(ctx, cfg)
 	expected = &Data{
-		ctx: ctx,
+		Ctx: actual.Ctx,
 		Sources: map[string]*Source{
 			"foo": {
 				Alias: "foo",
@@ -399,18 +403,18 @@ func TestFromConfig(t *testing.T) {
 			"bar": {
 				Alias: "bar",
 				URL:   mustParseURL("http://bar.com"),
-				header: http.Header{
+				Header: http.Header{
 					"Foo": []string{"bar"},
 				},
 			},
 		},
-		extraHeaders: map[string]http.Header{
+		ExtraHeaders: map[string]http.Header{
 			"baz": {
 				"Foo": []string{"bar"},
 			},
 		},
 	}
-	assert.EqualValues(t, expected, FromConfig(ctx, cfg))
+	assert.EqualValues(t, expected, actual)
 }
 
 func TestListDatasources(t *testing.T) {