diff --git a/libs/cmdio/io.go b/libs/cmdio/io.go index c0e9e868a..0d7ba5da8 100644 --- a/libs/cmdio/io.go +++ b/libs/cmdio/io.go @@ -285,3 +285,13 @@ func fromContext(ctx context.Context) *cmdIO { } return io } + +func MockContext(ctx context.Context) context.Context { + return InContext(ctx, &cmdIO{ + interactive: false, + outputFormat: flags.OutputText, + in: io.NopCloser(strings.NewReader("")), + out: io.Discard, + err: io.Discard, + }) +} diff --git a/libs/template/resolve.go b/libs/template/resolver.go similarity index 97% rename from libs/template/resolve.go rename to libs/template/resolver.go index 2c3f021e2..8d1b26d8d 100644 --- a/libs/template/resolve.go +++ b/libs/template/resolver.go @@ -62,7 +62,6 @@ func (r Resolver) Resolve(ctx context.Context) (*Template, error) { // User should not directly select "custom" and instead should provide the // file path or the Git URL for the template directly. - // Custom is just for internal representation purposes. if templateName == Custom { return nil, ErrCustomSelected } diff --git a/libs/template/resolver_test.go b/libs/template/resolver_test.go new file mode 100644 index 000000000..9fde4e9e9 --- /dev/null +++ b/libs/template/resolver_test.go @@ -0,0 +1,87 @@ +package template + +import ( + "context" + "testing" + + "github.com/databricks/cli/libs/cmdio" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestTemplateResolverBothTagAndBranch(t *testing.T) { + r := Resolver{ + Tag: "tag", + Branch: "branch", + } + + _, err := r.Resolve(context.Background()) + assert.EqualError(t, err, "only one of --tag or --branch can be specified") +} + +func TestTemplateResolverErrorsWhenPromptingIsNotSupported(t *testing.T) { + r := Resolver{} + ctx := cmdio.MockContext(context.Background()) + + _, err := r.Resolve(ctx) + assert.EqualError(t, err, "prompting is not supported. Please specify the path, name or URL of the template to use") +} + +func TestTemplateResolverErrorWhenUserSelectsCustom(t *testing.T) { + r := Resolver{ + TemplatePathOrUrl: "custom", + } + + _, err := r.Resolve(context.Background()) + assert.EqualError(t, err, "custom template selected") +} + +func TestTemplateResolverForDefaultTemplates(t *testing.T) { + for _, name := range []string{ + "default-python", + "default-sql", + "dbt-sql", + } { + r := Resolver{ + TemplatePathOrUrl: name, + } + + tmpl, err := r.Resolve(context.Background()) + require.NoError(t, err) + + assert.Equal(t, &builtinReader{name: name}, tmpl.Reader) + assert.IsType(t, &writerWithTelemetry{}, tmpl.Writer) + } + + r := Resolver{ + TemplatePathOrUrl: "mlops-stacks", + ConfigFile: "/config/file", + } + + tmpl, err := r.Resolve(context.Background()) + require.NoError(t, err) + + // Assert reader and writer configuration + assert.Equal(t, "https://github.com/databricks/mlops-stacks", tmpl.Reader.(*gitReader).gitUrl) + assert.Equal(t, "/config/file", tmpl.Writer.(*writerWithTelemetry).configPath) +} + +func TestTemplateResolverForCustomTemplate(t *testing.T) { + r := Resolver{ + TemplatePathOrUrl: "https://www.example.com/abc", + Tag: "tag", + TemplateDir: "/template/dir", + ConfigFile: "/config/file", + } + + tmpl, err := r.Resolve(context.Background()) + require.NoError(t, err) + + // Assert reader configuration + assert.Equal(t, "https://www.example.com/abc", tmpl.Reader.(*gitReader).gitUrl) + assert.Equal(t, "tag", tmpl.Reader.(*gitReader).ref) + assert.Equal(t, "/template/dir", tmpl.Reader.(*gitReader).templateDir) + + // Assert writer configuration + assert.Equal(t, "/config/file", tmpl.Writer.(*defaultWriter).configPath) +} diff --git a/libs/template/template.go b/libs/template/template.go index 46bdef57a..31c6e2814 100644 --- a/libs/template/template.go +++ b/libs/template/template.go @@ -27,7 +27,11 @@ const ( DbtSql TemplateName = "dbt-sql" MlopsStacks TemplateName = "mlops-stacks" DefaultPydabs TemplateName = "default-pydabs" - Custom TemplateName = "custom" + + // Custom represents any template that is not one of the above default + // templates. It's a catch for any custom templates that customers provide + // as a path or URL. + Custom TemplateName = "custom" ) var allTemplates = []Template{ @@ -98,7 +102,7 @@ func options() []cmdio.Tuple { func SelectTemplate(ctx context.Context) (TemplateName, error) { if !cmdio.IsPromptSupported(ctx) { - return "", fmt.Errorf("please specify a template") + return "", fmt.Errorf("prompting is not supported. Please specify the path, name or URL of the template to use") } description, err := cmdio.SelectOrdered(ctx, options(), "Template to use") if err != nil { diff --git a/libs/template/template_test.go b/libs/template/template_test.go index 73d818dfe..cfcdd6251 100644 --- a/libs/template/template_test.go +++ b/libs/template/template_test.go @@ -45,3 +45,18 @@ func TestBundleInitRepoName(t *testing.T) { assert.Equal(t, "invalid-url", repoName("invalid-url")) assert.Equal(t, "www.github.com", repoName("https://www.github.com")) } + +func TestTemplateTelemetryIsCapturedForAllDefaultTemplates(t *testing.T) { + for _, tmpl := range allTemplates { + w := tmpl.Writer + + if tmpl.name == Custom { + // Assert telemetry is not captured for user templates. + assert.IsType(t, &defaultWriter{}, w) + } else { + // Assert telemetry is captured for all other templates, i.e. templates + // owned by databricks. + assert.IsType(t, &writerWithTelemetry{}, w) + } + } +}