Fix --stream combined with --format for jsonnet

.drone.yml

@@ -3,6 +3,11 @@ type: docker
 name: default
 
 steps:
+- name: test
+  image: golang:1.16
+  commands:
+  - go test ./...
+
 - name: build
   image: golang:1.16
   commands:

drone/jsonnet/jsonnet.go

@@ -60,17 +60,30 @@ var Command = cli.Command{
 }
 
 func generate(c *cli.Context) error {
-	source := c.String("source")
+	result, err := convert(c.String("source"), c.Bool("string"), c.Bool("format"), c.Bool("stream"), c.StringSlice("extVar"))
+	if err != nil {
+		return err
+	}
+
+	// the user can optionally write the yaml to stdout. This is useful for debugging purposes without mutating an existing file.
+	if c.Bool("stdout") {
+		io.WriteString(os.Stdout, result)
+		return nil
+	}
+
 	target := c.String("target")
+	return ioutil.WriteFile(target, []byte(result), 0644)
+}
 
+func convert(source string, stringOutput bool, format bool, stream bool, vars []string) (string, error) {
 	data, err := ioutil.ReadFile(source)
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	vm := jsonnet.MakeVM()
 	vm.MaxStack = 500
-	vm.StringOutput = c.Bool("string")
+	vm.StringOutput = stringOutput
 	vm.ErrorFormatter.SetMaxStackTraceSize(20)
 	vm.ErrorFormatter.SetColorFormatter(
 		color.New(color.FgRed).Fprintf,
@@ -80,49 +93,55 @@ func generate(c *cli.Context) error {
 	RegisterNativeFuncs(vm)
 
 	// extVars
-	vars := c.StringSlice("extVar")
 	for _, v := range vars {
 		name, value, err := getVarVal(v)
 		if err != nil {
-			return err
+			return "", err
 		}
 		vm.ExtVar(name, value)
 	}
 
+	formatDoc := func(doc []byte) ([]byte, error) {
+		// enable yaml output
+		if format {
+			formatted, yErr := yaml.JSONToYAML(doc)
+			if yErr != nil {
+				return nil, fmt.Errorf("failed to convert to YAML: %v", yErr)
+			}
+			return formatted, nil
+		}
+		return doc, nil
+	}
+
 	buf := new(bytes.Buffer)
-	if c.Bool("stream") {
+	if stream {
 		docs, err := vm.EvaluateSnippetStream(source, string(data))
 		if err != nil {
-			return err
+			return "", err
 		}
 		for _, doc := range docs {
+			formatted, err := formatDoc([]byte(doc))
+			if err != nil {
+				return "", err
+			}
+
 			buf.WriteString("---")
 			buf.WriteString("\n")
-			buf.WriteString(doc)
+			buf.Write(formatted)
 		}
 	} else {
 		result, err := vm.EvaluateSnippet(source, string(data))
 		if err != nil {
-			return err
+			return "", err
 		}
-		buf.WriteString(result)
-	}
-	// enable yaml output
-	if c.Bool("format") {
-		formatted, yErr := yaml.JSONToYAML(buf.Bytes())
-		if yErr != nil {
-			return fmt.Errorf("failed to convert to YAML: %v", yErr)
+		formatted, err := formatDoc([]byte(result))
+		if err != nil {
+			return "", err
 		}
-		buf.Reset()
 		buf.Write(formatted)
 	}
-	// the user can optionally write the yaml to stdout. This is useful for debugging purposes without mutating an existing file.
-	if c.Bool("stdout") {
-		io.Copy(os.Stdout, buf)
-		return nil
-	}
 
-	return ioutil.WriteFile(target, buf.Bytes(), 0644)
+	return buf.String(), nil
 }
 
 // https://github.com/google/go-jsonnet/blob/master/cmd/jsonnet/cmd.go#L149

drone/jsonnet/jsonnet_test.go

@@ -0,0 +1,37 @@
+package jsonnet
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestConvert(t *testing.T) {
+	testcases := []struct {
+		name                         string
+		jsonnetFile, yamlFile        string
+		stringOutput, format, stream bool
+		extVars                      []string
+	}{
+		{
+			name:        "Stream + Format",
+			jsonnetFile: "stream_format.jsonnet",
+			yamlFile:    "stream_format.yaml",
+			format:      true, stream: true,
+		},
+	}
+
+	for _, tc := range testcases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			expected, err := os.ReadFile(filepath.Join("./testdata", tc.yamlFile))
+			assert.NoError(t, err)
+
+			result, err := convert(filepath.Join("./testdata", tc.jsonnetFile), tc.stringOutput, tc.format, tc.stream, tc.extVars)
+			assert.NoError(t, err)
+			assert.Equal(t, string(expected), result)
+		})
+	}
+}

drone/jsonnet/testdata/stream_format.jsonnet

@@ -0,0 +1,27 @@
+local pipeline(name) =
+  {
+    kind: 'pipeline',
+    type: 'docker',
+    name: name,
+    platform: {
+      os: 'linux',
+      arch: 'amd64',
+    },
+    steps: [
+      {
+        name: 'test',
+        image: 'golang:1.16',
+        commands: ['go test ./...'],
+      },
+      {
+        name: 'build',
+        image: 'golang:1.16',
+        commands: ['go build ./...'],
+      },
+    ],
+  };
+
+[
+  pipeline('first'),
+  pipeline('second'),
+]

drone/jsonnet/testdata/stream_format.yaml

@@ -0,0 +1,32 @@
+---
+kind: pipeline
+name: first
+platform:
+  arch: amd64
+  os: linux
+steps:
+- commands:
+  - go test ./...
+  image: golang:1.16
+  name: test
+- commands:
+  - go build ./...
+  image: golang:1.16
+  name: build
+type: docker
+---
+kind: pipeline
+name: second
+platform:
+  arch: amd64
+  os: linux
+steps:
+- commands:
+  - go test ./...
+  image: golang:1.16
+  name: test
+- commands:
+  - go build ./...
+  image: golang:1.16
+  name: build
+type: docker

go.mod

@@ -20,6 +20,7 @@ require (
 	github.com/mattn/go-colorable v0.1.4
 	github.com/mattn/go-isatty v0.0.11
 	github.com/pkg/browser v0.0.0-20180916011732-0a3d74bf9ce4
+	github.com/stretchr/testify v1.4.0
 	github.com/urfave/cli v1.20.0
 	go.starlark.net v0.0.0-20201118183435-e55f603d8c79
 	golang.org/x/net v0.0.0-20190603091049-60506f45cf65