From 2e7a37890cb87cc2e0fcd6f349b20f78f5f6f203 Mon Sep 17 00:00:00 2001
From: Matt Keeler <mjkeeler7@gmail.com>
Date: Mon, 21 May 2018 15:32:25 -0400
Subject: [PATCH] Move data source loading into a command helpers function

---
 command/helpers/helpers.go | 42 ++++++++++++++++++++++++++++++++++++++
 command/kv/put/kv_put.go   | 36 +++++---------------------------
 2 files changed, 47 insertions(+), 31 deletions(-)
 create mode 100644 command/helpers/helpers.go

diff --git a/command/helpers/helpers.go b/command/helpers/helpers.go
new file mode 100644
index 0000000000..6ad7ed2b7d
--- /dev/null
+++ b/command/helpers/helpers.go
@@ -0,0 +1,42 @@
+package helpers
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+)
+
+func LoadDataSource(data string, testStdin io.Reader) (string, error) {
+	var stdin io.Reader = os.Stdin
+	if testStdin != nil {
+		stdin = testStdin
+	}
+
+	// Handle empty quoted shell parameters
+	if len(data) == 0 {
+		return "", nil
+	}
+
+	switch data[0] {
+	case '@':
+		data, err := ioutil.ReadFile(data[1:])
+		if err != nil {
+			return "", fmt.Errorf("Failed to read file: %s", err)
+		} else {
+			return string(data), nil
+		}
+	case '-':
+		if len(data) > 1 {
+			return data, nil
+		}
+		var b bytes.Buffer
+		if _, err := io.Copy(&b, stdin); err != nil {
+			return "", fmt.Errorf("Failed to read stdin: %s", err)
+		}
+		return b.String(), nil
+	default:
+		return data, nil
+	}
+}
diff --git a/command/kv/put/kv_put.go b/command/kv/put/kv_put.go
index 98be3c001e..abe51a5382 100644
--- a/command/kv/put/kv_put.go
+++ b/command/kv/put/kv_put.go
@@ -1,16 +1,14 @@
 package put
 
 import (
-	"bytes"
 	"encoding/base64"
 	"flag"
 	"fmt"
 	"io"
-	"io/ioutil"
-	"os"
 
 	"github.com/hashicorp/consul/api"
 	"github.com/hashicorp/consul/command/flags"
+	"github.com/hashicorp/consul/command/helpers"
 	"github.com/mitchellh/cli"
 )
 
@@ -173,11 +171,6 @@ func (c *cmd) Run(args []string) int {
 }
 
 func (c *cmd) dataFromArgs(args []string) (string, string, error) {
-	var stdin io.Reader = os.Stdin
-	if c.testStdin != nil {
-		stdin = c.testStdin
-	}
-
 	switch len(args) {
 	case 0:
 		return "", "", fmt.Errorf("Missing KEY argument")
@@ -189,30 +182,11 @@ func (c *cmd) dataFromArgs(args []string) (string, string, error) {
 	}
 
 	key := args[0]
-	data := args[1]
+	data, err := helpers.LoadDataSource(args[1], c.testStdin)
 
-	// Handle empty quoted shell parameters
-	if len(data) == 0 {
-		return key, "", nil
-	}
-
-	switch data[0] {
-	case '@':
-		data, err := ioutil.ReadFile(data[1:])
-		if err != nil {
-			return "", "", fmt.Errorf("Failed to read file: %s", err)
-		}
-		return key, string(data), nil
-	case '-':
-		if len(data) > 1 {
-			return key, data, nil
-		}
-		var b bytes.Buffer
-		if _, err := io.Copy(&b, stdin); err != nil {
-			return "", "", fmt.Errorf("Failed to read stdin: %s", err)
-		}
-		return key, b.String(), nil
-	default:
+	if err != nil {
+		return "", "", err
+	} else {
 		return key, data, nil
 	}
 }