mirror of https://github.com/databricks/cli.git
90 lines
2.3 KiB
Go
90 lines
2.3 KiB
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.FrontendLogEntry{
|
|
{
|
|
DatabricksCliLog: telemetry.DatabricksCliLog{
|
|
CliTestEvent: &events.CliTestEvent{
|
|
Name: events.DummyCliEnumValue1,
|
|
},
|
|
},
|
|
},
|
|
{
|
|
DatabricksCliLog: telemetry.DatabricksCliLog{
|
|
BundleInitEvent: &events.BundleInitEvent{
|
|
Uuid: uuid.New().String(),
|
|
TemplateName: "abc",
|
|
TemplateEnumArgs: map[string]string{
|
|
"a": "b",
|
|
"c": "d",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
assert.Equal(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 {
|
|
err := telemetry.Log(ctx, event)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
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)
|
|
}
|