more tests and cleanup

This commit is contained in:
Shreyas Goenka 2024-12-27 14:32:42 +05:30
parent 598e569957
commit 4563397edc
No known key found for this signature in database
GPG Key ID: 92A07DF49CCB0622
3 changed files with 14 additions and 6 deletions

View File

@ -257,10 +257,11 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
} }
nt := getNativeTemplateByName(templatePath) nt := getNativeTemplateByName(templatePath)
templateName := "custom" var templateName string
isTemplateDatabricksOwned := false isTemplateDatabricksOwned := false
if nt != nil { 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 templateName = templatePath
// if we have a Git URL for the native template, expand templatePath // if we have a Git URL for the native template, expand templatePath

View File

@ -204,7 +204,7 @@ func TestBundleInitTelemetryForCustomTemplates(t *testing.T) {
err := os.Mkdir(filepath.Join(tmpDir1, "template"), 0o755) err := os.Mkdir(filepath.Join(tmpDir1, "template"), 0o755)
require.NoError(t, err) 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) require.NoError(t, err)
err = os.WriteFile(filepath.Join(tmpDir1, "databricks_template_schema.json"), []byte(` 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 event := logs[0].Entry.DatabricksCliLog.BundleInitEvent
assert.Equal(t, event.TemplateName, "custom") assert.Equal(t, event.TemplateName, "custom")
assert.Nil(t, event.TemplateEnumArgs) 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) { func TestBundleInitHelpers(t *testing.T) {

View File

@ -25,10 +25,10 @@ type TemplateOpts struct {
TemplateFS fs.FS TemplateFS fs.FS
// filer to use for writing the initialized template // filer to use for writing the initialized template
OutputFiler filer.Filer 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 IsDatabricksOwned bool
// Name of the template. For non-Databricks owned templates, this is set to // Name of the template. For non-Databricks owned templates, this is set to
// custom. // "custom".
Name string Name string
} }
@ -108,10 +108,12 @@ func (t *Template) printSuccessMessage(ctx context.Context) error {
func (t *Template) logTelemetry(ctx context.Context) error { func (t *Template) logTelemetry(ctx context.Context) error {
// Only log telemetry input for Databricks owned templates. This is to prevent // 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{} templateEnumArgs := map[string]string{}
if t.IsDatabricksOwned { if t.IsDatabricksOwned {
templateEnumArgs = t.config.enumValues() templateEnumArgs = t.config.enumValues()
} else {
t.Name = "custom"
} }
event := telemetry.DatabricksCliLog{ event := telemetry.DatabricksCliLog{