From 4563397edc81d8e3a5294f6f03bd45f16fb8f53e Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Fri, 27 Dec 2024 14:32:42 +0530
Subject: [PATCH] more tests and cleanup

---
 cmd/bundle/init.go              | 5 +++--
 integration/bundle/init_test.go | 7 ++++++-
 libs/template/materialize.go    | 8 +++++---
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/cmd/bundle/init.go b/cmd/bundle/init.go
index 1881c4ca..2fff63c1 100644
--- a/cmd/bundle/init.go
+++ b/cmd/bundle/init.go
@@ -257,10 +257,11 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 		}
 
 		nt := getNativeTemplateByName(templatePath)
-		templateName := "custom"
+		var templateName string
 		isTemplateDatabricksOwned := false
 		if nt != nil {
-			// If the template is a native template, set it's value.
+			// If the template is a native template, templatePath is the name of the template.
+			// Eg: templatePath = "default-python".
 			templateName = templatePath
 
 			// if we have a Git URL for the native template, expand templatePath
diff --git a/integration/bundle/init_test.go b/integration/bundle/init_test.go
index bec64602..537d8bf8 100644
--- a/integration/bundle/init_test.go
+++ b/integration/bundle/init_test.go
@@ -204,7 +204,7 @@ func TestBundleInitTelemetryForCustomTemplates(t *testing.T) {
 
 	err := os.Mkdir(filepath.Join(tmpDir1, "template"), 0o755)
 	require.NoError(t, err)
-	err = os.WriteFile(filepath.Join(tmpDir1, "template", "foo.txt.tmpl"), []byte("doesnotmatter"), 0o644)
+	err = os.WriteFile(filepath.Join(tmpDir1, "template", "foo.txt.tmpl"), []byte("{{bundle_uuid}}"), 0o644)
 	require.NoError(t, err)
 	err = os.WriteFile(filepath.Join(tmpDir1, "databricks_template_schema.json"), []byte(`
 {
@@ -247,6 +247,11 @@ func TestBundleInitTelemetryForCustomTemplates(t *testing.T) {
 	event := logs[0].Entry.DatabricksCliLog.BundleInitEvent
 	assert.Equal(t, event.TemplateName, "custom")
 	assert.Nil(t, event.TemplateEnumArgs)
+
+	// Ensure that the UUID returned by the `bundle_uuid` helper is the same UUID
+	// that's logged in the telemetry event.
+	fileC := testutil.ReadFile(t, filepath.Join(tmpDir2, "foo.txt"))
+	assert.Equal(t, event.Uuid, fileC)
 }
 
 func TestBundleInitHelpers(t *testing.T) {
diff --git a/libs/template/materialize.go b/libs/template/materialize.go
index 3ef75773..a2bae355 100644
--- a/libs/template/materialize.go
+++ b/libs/template/materialize.go
@@ -25,10 +25,10 @@ type TemplateOpts struct {
 	TemplateFS fs.FS
 	// filer to use for writing the initialized template
 	OutputFiler filer.Filer
-	// If true, we'll include the enum template args in the telemetry payload..
+	// If true, we'll include the enum template args in the telemetry payload.
 	IsDatabricksOwned bool
 	// Name of the template. For non-Databricks owned templates, this is set to
-	// custom.
+	// "custom".
 	Name string
 }
 
@@ -108,10 +108,12 @@ func (t *Template) printSuccessMessage(ctx context.Context) error {
 
 func (t *Template) logTelemetry(ctx context.Context) error {
 	// Only log telemetry input for Databricks owned templates. This is to prevent
-	// accidentally collecting PUII from custom user templates.
+	// accidentally collecting PII from custom user templates.
 	templateEnumArgs := map[string]string{}
 	if t.IsDatabricksOwned {
 		templateEnumArgs = t.config.enumValues()
+	} else {
+		t.Name = "custom"
 	}
 
 	event := telemetry.DatabricksCliLog{