databricks-cli/libs/cmdio/render_test.go

package cmdio

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"testing"

	"github.com/databricks/cli/libs/flags"
	"github.com/databricks/databricks-sdk-go/listing"
	"github.com/databricks/databricks-sdk-go/service/provisioning"
	"github.com/stretchr/testify/assert"
)

type testCase struct {
	name           string
	v              any
	outputFormat   flags.Output
	headerTemplate string
	template       string
	expected       string
	errMessage     string
}

var dummyWorkspace1 = provisioning.Workspace{
	WorkspaceId:   123,
	WorkspaceName: "abc",
}

var dummyWorkspace2 = provisioning.Workspace{
	WorkspaceId:   456,
	WorkspaceName: "def",
}

type dummyIterator struct {
	items []*provisioning.Workspace
}

func (d *dummyIterator) HasNext(_ context.Context) bool {
	return len(d.items) > 0
}

func (d *dummyIterator) Next(ctx context.Context) (*provisioning.Workspace, error) {
	if !d.HasNext(ctx) {
		return nil, errors.New("no more items")
	}
	item := d.items[0]
	d.items = d.items[1:]
	return item, nil
}

func makeWorkspaces(count int) []*provisioning.Workspace {
	res := make([]*provisioning.Workspace, 0, count)
	next := []*provisioning.Workspace{&dummyWorkspace1, &dummyWorkspace2}
	for i := 0; i < count; i++ {
		n := next[0]
		next = append(next[1:], n)
		res = append(res, n)
	}
	return res
}

func makeIterator(count int) listing.Iterator[*provisioning.Workspace] {
	items := make([]*provisioning.Workspace, 0, count)
	items = append(items, makeWorkspaces(count)...)
	return &dummyIterator{
		items: items,
	}
}

func makeBigOutput(count int) string {
	res := bytes.Buffer{}
	for _, ws := range makeWorkspaces(count) {
		res.Write([]byte(fmt.Sprintf("%d  %s\n", ws.WorkspaceId, ws.WorkspaceName)))
	}
	return res.String()
}

func must[T any](a T, e error) T {
	if e != nil {
		panic(e)
	}
	return a
}

var testCases = []testCase{
	{
		name:           "Workspace with header and template",
		v:              dummyWorkspace1,
		outputFormat:   flags.OutputText,
		headerTemplate: "id\tname",
		template:       "{{.WorkspaceId}}\t{{.WorkspaceName}}",
		expected: `id   name
123  abc`,
	},
	{
		name:         "Workspace with no header and template",
		v:            dummyWorkspace1,
		outputFormat: flags.OutputText,
		template:     "{{.WorkspaceId}}\t{{.WorkspaceName}}",
		expected:     `123  abc`,
	},
	{
		name:         "Workspace with no header and no template",
		v:            dummyWorkspace1,
		outputFormat: flags.OutputText,
		expected: `{
  "workspace_id":123,
  "workspace_name":"abc"
}
`,
	},
	{
		name:           "Workspace Iterator with header and template",
		v:              makeIterator(2),
		outputFormat:   flags.OutputText,
		headerTemplate: "id\tname",
		template:       "{{range .}}{{.WorkspaceId}}\t{{.WorkspaceName}}\n{{end}}",
		expected: `id   name
123  abc
456  def
`,
	},
	{
		name:         "Workspace Iterator with no header and template",
		v:            makeIterator(2),
		outputFormat: flags.OutputText,
		template:     "{{range .}}{{.WorkspaceId}}\t{{.WorkspaceName}}\n{{end}}",
		expected: `123  abc
456  def
`,
	},
	{
		name:         "Workspace Iterator with no header and no template",
		v:            makeIterator(2),
		outputFormat: flags.OutputText,
		expected:     string(must(json.MarshalIndent(makeWorkspaces(2), "", "  "))) + "\n",
	},
	{
		name:           "Big Workspace Iterator with template",
		v:              makeIterator(234),
		outputFormat:   flags.OutputText,
		headerTemplate: "id\tname",
		template:       "{{range .}}{{.WorkspaceId}}\t{{.WorkspaceName}}\n{{end}}",
		expected:       "id   name\n" + makeBigOutput(234),
	},
	{
		name:         "Big Workspace Iterator with no template",
		v:            makeIterator(234),
		outputFormat: flags.OutputText,
		expected:     string(must(json.MarshalIndent(makeWorkspaces(234), "", "  "))) + "\n",
	},
	{
		name:         "io.Reader",
		v:            strings.NewReader("a test"),
		outputFormat: flags.OutputText,
		expected:     "a test",
	},
	{
		name:         "io.Reader",
		v:            strings.NewReader("a test"),
		outputFormat: flags.OutputJSON,
		errMessage:   "json output not supported",
	},
}

func TestRender(t *testing.T) {
	for _, c := range testCases {
		t.Run(c.name, func(t *testing.T) {
			output := &bytes.Buffer{}
			cmdIO := NewIO(c.outputFormat, nil, output, output, c.headerTemplate, c.template)
			ctx := InContext(context.Background(), cmdIO)
			var err error
			if vv, ok := c.v.(listing.Iterator[*provisioning.Workspace]); ok {
				err = RenderIterator(ctx, vv)
			} else {
				err = Render(ctx, c.v)
			}
			if c.errMessage != "" {
				assert.ErrorContains(t, err, c.errMessage)
			} else {
				assert.NoError(t, err)
				assert.Equal(t, c.expected, output.String())
			}
		})
	}
}