diff --git a/integration/bundle/init_test.go b/integration/bundle/init_test.go index 3826f5543..f5c263ca3 100644 --- a/integration/bundle/init_test.go +++ b/integration/bundle/init_test.go @@ -15,7 +15,6 @@ import ( "github.com/databricks/cli/internal/testcli" "github.com/databricks/cli/internal/testutil" "github.com/databricks/cli/libs/iamutil" - "github.com/databricks/cli/libs/telemetry" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -43,9 +42,6 @@ func TestBundleInitOnMlopsStacks(t *testing.T) { ctx, wt := acc.WorkspaceTest(t) w := wt.W - // Use mock logger to introspect the telemetry payload. - ctx = telemetry.WithMockLogger(ctx) - tmpDir1 := t.TempDir() tmpDir2 := t.TempDir() @@ -68,28 +64,6 @@ func TestBundleInitOnMlopsStacks(t *testing.T) { assert.NoFileExists(t, filepath.Join(tmpDir2, "repo_name", projectName, "README.md")) testcli.RequireSuccessfulRun(t, ctx, "bundle", "init", "mlops-stacks", "--output-dir", tmpDir2, "--config-file", filepath.Join(tmpDir1, "config.json")) - // Assert the telemetry payload is correctly logged. - tlmyEvents := telemetry.Introspect(ctx) - require.Len(t, telemetry.Introspect(ctx), 1) - event := tlmyEvents[0].BundleInitEvent - assert.Equal(t, "mlops-stacks", event.TemplateName) - - get := func(key string) string { - for _, v := range event.TemplateEnumArgs { - if v.Key == key { - return v.Value - } - } - return "" - } - - // Enum values should be present in the telemetry payload. - assert.Equal(t, "no", get("input_include_models_in_unity_catalog")) - assert.Equal(t, strings.ToLower(env), get("input_cloud")) - // Freeform strings should not be present in the telemetry payload. - assert.Equal(t, "", get("input_project_name")) - assert.Equal(t, "", get("input_root_dir")) - // Assert that the README.md file was created contents := testutil.ReadFile(t, filepath.Join(tmpDir2, "repo_name", projectName, "README.md")) assert.Contains(t, contents, fmt.Sprintf("# %s", projectName)) @@ -125,156 +99,6 @@ func TestBundleInitOnMlopsStacks(t *testing.T) { assert.Contains(t, job.Settings.Name, fmt.Sprintf("dev-%s-batch-inference-job", projectName)) } -func TestBundleInitTelemetryForDefaultTemplates(t *testing.T) { - projectName := testutil.RandomName("name_") - - tcases := []struct { - name string - args map[string]string - expectedArgs map[string]string - }{ - { - name: "dbt-sql", - args: map[string]string{ - "project_name": fmt.Sprintf("dbt-sql-%s", projectName), - "http_path": "/sql/1.0/warehouses/id", - "default_catalog": "abcd", - "personal_schemas": "yes, use a schema based on the current user name during development", - }, - expectedArgs: map[string]string{ - "personal_schemas": "yes, use a schema based on the current user name during development", - }, - }, - { - name: "default-python", - args: map[string]string{ - "project_name": fmt.Sprintf("default_python_%s", projectName), - "include_notebook": "yes", - "include_dlt": "yes", - "include_python": "no", - }, - expectedArgs: map[string]string{ - "include_notebook": "yes", - "include_dlt": "yes", - "include_python": "no", - }, - }, - { - name: "default-sql", - args: map[string]string{ - "project_name": fmt.Sprintf("sql_project_%s", projectName), - "http_path": "/sql/1.0/warehouses/id", - "default_catalog": "abcd", - "personal_schemas": "yes, automatically use a schema based on the current user name during development", - }, - expectedArgs: map[string]string{ - "personal_schemas": "yes, automatically use a schema based on the current user name during development", - }, - }, - } - - for _, tc := range tcases { - ctx, _ := acc.WorkspaceTest(t) - - // Use mock logger to introspect the telemetry payload. - ctx = telemetry.WithMockLogger(ctx) - - tmpDir1 := t.TempDir() - tmpDir2 := t.TempDir() - - // Create a config file with the project name and root dir - initConfig := tc.args - b, err := json.Marshal(initConfig) - require.NoError(t, err) - err = os.WriteFile(filepath.Join(tmpDir1, "config.json"), b, 0o644) - require.NoError(t, err) - - // Run bundle init - assert.NoDirExists(t, filepath.Join(tmpDir2, tc.args["project_name"])) - testcli.RequireSuccessfulRun(t, ctx, "bundle", "init", tc.name, "--output-dir", tmpDir2, "--config-file", filepath.Join(tmpDir1, "config.json")) - assert.DirExists(t, filepath.Join(tmpDir2, tc.args["project_name"])) - - // Assert the telemetry payload is correctly logged. - logs := telemetry.Introspect(ctx) - require.Len(t, logs, 1) - event := logs[0].BundleInitEvent - assert.Equal(t, event.TemplateName, tc.name) - - get := func(key string) string { - for _, v := range event.TemplateEnumArgs { - if v.Key == key { - return v.Value - } - } - return "" - } - - // Assert the template enum args are correctly logged. - assert.Len(t, event.TemplateEnumArgs, len(tc.expectedArgs)) - for k, v := range tc.expectedArgs { - assert.Equal(t, get(k), v) - } - } -} - -func TestBundleInitTelemetryForCustomTemplates(t *testing.T) { - ctx, _ := acc.WorkspaceTest(t) - - tmpDir1 := t.TempDir() - tmpDir2 := t.TempDir() - tmpDir3 := t.TempDir() - - err := os.Mkdir(filepath.Join(tmpDir1, "template"), 0o755) - require.NoError(t, err) - 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(` -{ - "properties": { - "a": { - "description": "whatever", - "type": "string" - }, - "b": { - "description": "whatever", - "type": "string", - "enum": ["yes", "no"] - } - } -} -`), 0o644) - require.NoError(t, err) - - // Create a config file with the project name and root dir - initConfig := map[string]string{ - "a": "v1", - "b": "yes", - } - b, err := json.Marshal(initConfig) - require.NoError(t, err) - err = os.WriteFile(filepath.Join(tmpDir3, "config.json"), b, 0o644) - require.NoError(t, err) - - // Use mock logger to introspect the telemetry payload. - ctx = telemetry.WithMockLogger(ctx) - - // Run bundle init. - 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. - tlmyEvents := telemetry.Introspect(ctx) - require.Len(t, len(tlmyEvents), 1) - event := tlmyEvents[0].BundleInitEvent - assert.Equal(t, "custom", event.TemplateName) - assert.Empty(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) { ctx, wt := acc.WorkspaceTest(t) w := wt.W