From 2eb3344edba323a5b12adf495165e845c86d283b Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Fri, 27 Dec 2024 13:42:06 +0530
Subject: [PATCH] fix tests

---
 integration/bundle/helpers_test.go | 2 ++
 integration/bundle/init_test.go    | 2 +-
 libs/telemetry/logger.go           | 4 ++--
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/integration/bundle/helpers_test.go b/integration/bundle/helpers_test.go
index 84c8364dd..a6edfed70 100644
--- a/integration/bundle/helpers_test.go
+++ b/integration/bundle/helpers_test.go
@@ -19,6 +19,7 @@ import (
 	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/cli/libs/folders"
+	"github.com/databricks/cli/libs/telemetry"
 	"github.com/databricks/cli/libs/template"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/stretchr/testify/require"
@@ -28,6 +29,7 @@ const defaultSparkVersion = "13.3.x-snapshot-scala2.12"
 
 func initTestTemplate(t testutil.TestingT, ctx context.Context, templateName string, config map[string]any) string {
 	bundleRoot := t.TempDir()
+	ctx = telemetry.ContextWithLogger(ctx)
 	return initTestTemplateWithBundleRoot(t, ctx, templateName, config, bundleRoot)
 }
 
diff --git a/integration/bundle/init_test.go b/integration/bundle/init_test.go
index f80f6f8f3..bec646021 100644
--- a/integration/bundle/init_test.go
+++ b/integration/bundle/init_test.go
@@ -240,7 +240,7 @@ func TestBundleInitTelemetryForCustomTemplates(t *testing.T) {
 	testcli.RequireSuccessfulRun(t, ctx, "bundle", "init", tmpDir1, "--output-dir", tmpDir2, "--config-file", filepath.Join(tmpDir3, "config.json"))
 
 	// Assert the telemetry payload is correctly logged. For custom templates we should
-	//
+	// never set template_enum_args.
 	logs, err := telemetry.GetLogs(ctx)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(logs))
diff --git a/libs/telemetry/logger.go b/libs/telemetry/logger.go
index 84e828b3e..a22abbd93 100644
--- a/libs/telemetry/logger.go
+++ b/libs/telemetry/logger.go
@@ -41,8 +41,8 @@ type logger struct {
 	protoLogs []string
 }
 
-// Only to be used in tests to introspect the telemetry logs that are queued
-// to be flushed.
+// This function is meant to be only to be used in tests to introspect the telemetry logs
+// that have been logged so far.
 func GetLogs(ctx context.Context) ([]FrontendLog, error) {
 	l := fromContext(ctx)
 	res := []FrontendLog{}