databricks-cli/integration/libs/telemetry/telemetry_test.go

package telemetry_test

import (
	"context"
	"net/http"
	"reflect"
	"testing"
	"time"

	"github.com/databricks/cli/integration/internal/acc"
	"github.com/databricks/cli/libs/telemetry"
	"github.com/databricks/cli/libs/telemetry/events"
	"github.com/databricks/databricks-sdk-go/client"
	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// Wrapper to capture the response from the API client since that's not directly
// accessible from the logger.
type apiClientWrapper struct {
	response  *telemetry.ResponseBody
	apiClient *client.DatabricksClient
}

func (wrapper *apiClientWrapper) Do(ctx context.Context, method, path string,
	headers map[string]string, request, response any,
	visitors ...func(*http.Request) error,
) error {
	err := wrapper.apiClient.Do(ctx, method, path, headers, request, response, visitors...)
	wrapper.response = response.(*telemetry.ResponseBody)
	return err
}

func TestTelemetryLogger(t *testing.T) {
	events := []telemetry.DatabricksCliLog{
		{
			CliTestEvent: &events.CliTestEvent{
				Name: events.DummyCliEnumValue1,
			},
		},
		{
			BundleInitEvent: &events.BundleInitEvent{
				Uuid:         uuid.New().String(),
				TemplateName: "abc",
				TemplateEnumArgs: []events.BundleInitTemplateEnumArg{
					{
						Key:   "a",
						Value: "b",
					},
					{
						Key:   "c",
						Value: "d",
					},
				},
			},
		},
	}

	assert.Len(t, reflect.TypeOf(telemetry.DatabricksCliLog{}).NumField(), len(events),
		"Number of events should match the number of fields in DatabricksCliLog. Please add a new event to this test.")

	ctx, w := acc.WorkspaceTest(t)
	ctx = telemetry.ContextWithLogger(ctx)

	// Extend the maximum wait time for the telemetry flush just for this test.
	telemetry.MaxAdditionalWaitTime = 1 * time.Hour
	t.Cleanup(func() {
		telemetry.MaxAdditionalWaitTime = 2 * time.Second
	})

	for _, event := range events {
		telemetry.Log(ctx, event)
	}

	apiClient, err := client.New(w.W.Config)
	require.NoError(t, err)

	// Flush the events.
	wrapper := &apiClientWrapper{
		apiClient: apiClient,
	}
	telemetry.Flush(ctx, wrapper)

	// Assert that the events were logged.
	assert.Equal(t, telemetry.ResponseBody{
		NumProtoSuccess: int64(len(events)),
		Errors:          []telemetry.LogError{},
	}, *wrapper.response)
}