From ecb977f4edb7774dcc0ddb33cfe4b578b5b4fac2 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Fri, 27 Dec 2024 12:19:42 +0530
Subject: [PATCH] cleanup code

---
 cmd/bundle/init.go           |  3 +--
 libs/telemetry/logger.go     | 14 +------------
 libs/template/config.go      | 10 ++++++++-
 libs/template/config_test.go | 39 ++++++++++++++++++++++++++++++++++++
 libs/template/materialize.go |  1 -
 5 files changed, 50 insertions(+), 17 deletions(-)

diff --git a/cmd/bundle/init.go b/cmd/bundle/init.go
index 9b945df6b..4b0040728 100644
--- a/cmd/bundle/init.go
+++ b/cmd/bundle/init.go
@@ -235,12 +235,11 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 			return nil
 		}
 
-		// If templatePath refers to a native template, store it's name in a new
-		// variable.
 		nt := getNativeTemplateByName(templatePath)
 		templateName := "custom"
 		isTemplateDatabricksOwned := false
 		if nt != nil {
+			// If the template is a native template, set it's value.
 			templateName = templatePath
 
 			// if we have a Git URL for the native template, expand templatePath
diff --git a/libs/telemetry/logger.go b/libs/telemetry/logger.go
index 844d56b84..ce652d675 100644
--- a/libs/telemetry/logger.go
+++ b/libs/telemetry/logger.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"slices"
 	"time"
 
 	"github.com/databricks/cli/libs/log"
@@ -61,19 +62,6 @@ func Flush(ctx context.Context, apiClient DatabricksApiClient) {
 		return
 	}
 
-	// We pass the API client as an arg to mock it in unit tests.
-	// TODO: Cleanup and remove this section.
-	// if apiClient == nil {
-	// 	var err error
-
-	// 	// Create API client to make the the telemetry API call.
-	// 	apiClient, err = client.New(root.WorkspaceClient(ctx).Config)
-	// 	if err != nil {
-	// 		log.Debugf(ctx, "error creating API client for telemetry: %v", err)
-	// 		return
-	// 	}
-	// }
-
 	resp := &ResponseBody{}
 	for {
 		select {
diff --git a/libs/template/config.go b/libs/template/config.go
index 251d3e9ae..34eee065c 100644
--- a/libs/template/config.go
+++ b/libs/template/config.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"io/fs"
+	"slices"
 
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/jsonschema"
@@ -279,10 +280,17 @@ func (c *config) validate() error {
 func (c *config) enumValues() map[string]string {
 	res := map[string]string{}
 	for k, p := range c.schema.Properties {
+		if p.Type != jsonschema.StringType {
+			continue
+		}
 		if p.Enum == nil {
 			continue
 		}
-		res[k] = c.values[k].(string)
+		v := c.values[k]
+
+		if slices.Contains(p.Enum, v) {
+			res[k] = v.(string)
+		}
 	}
 	return res
 }
diff --git a/libs/template/config_test.go b/libs/template/config_test.go
index a855019b6..fca17eba8 100644
--- a/libs/template/config_test.go
+++ b/libs/template/config_test.go
@@ -564,3 +564,42 @@ func TestPromptIsSkippedAnyOf(t *testing.T) {
 	assert.True(t, skip)
 	assert.Equal(t, "hello-world", c.values["xyz"])
 }
+
+func TestConfigEnumValues(t *testing.T) {
+	c := &config{
+		schema: &jsonschema.Schema{
+			Properties: map[string]*jsonschema.Schema{
+				"a": {
+					Type: jsonschema.StringType,
+				},
+				"b": {
+					Type: jsonschema.BooleanType,
+				},
+				"c": {
+					Type: jsonschema.StringType,
+					Enum: []any{"v1", "v2"},
+				},
+				"d": {
+					Type: jsonschema.StringType,
+					Enum: []any{"v3", "v4"},
+				},
+				"e": {
+					Type: jsonschema.StringType,
+					Enum: []any{"v5", "v6"},
+				},
+			},
+		},
+		values: map[string]any{
+			"a": "w1",
+			"b": false,
+			"c": "v1",
+			"d": "v3",
+			"e": "v7",
+		},
+	}
+
+	assert.Equal(t, map[string]string{
+		"c": "v1",
+		"d": "v3",
+	}, c.enumValues())
+}
diff --git a/libs/template/materialize.go b/libs/template/materialize.go
index 1590199d8..c0961b486 100644
--- a/libs/template/materialize.go
+++ b/libs/template/materialize.go
@@ -29,7 +29,6 @@ type TemplateOpts struct {
 	IsDatabricksOwned bool
 	// Name of the template. For non-Databricks owned templates, this is set to
 	// custom.
-	// TODO: move this enum to the template package.
 	Name string
 }