databricks-cli/libs/telemetry/logger_test.go

package telemetry

import (
	"context"
	"math/rand"
	"net/http"
	"testing"

	"github.com/databricks/cli/libs/telemetry/events"
	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
)

type mockDatabricksClient struct {
	numCalls int
	t        *testing.T
}

func (m *mockDatabricksClient) Do(ctx context.Context, method, path string, headers map[string]string, request, response any, visitors ...func(*http.Request) error) error {
	// Block until the fire channel is fired.
	m.numCalls++

	assertRequestPayload := func(reqb RequestBody) {
		expectedProtoLogs := []string{
			"{\"frontend_log_event_id\":\"0194fdc2-fa2f-4cc0-81d3-ff12045b73c8\",\"entry\":{\"databricks_cli_log\":{\"cli_test_event\":{\"name\":\"VALUE1\"}}}}",
			"{\"frontend_log_event_id\":\"6e4ff95f-f662-45ee-a82a-bdf44a2d0b75\",\"entry\":{\"databricks_cli_log\":{\"cli_test_event\":{\"name\":\"VALUE2\"}}}}",
			"{\"frontend_log_event_id\":\"fb180daf-48a7-4ee0-b10d-394651850fd4\",\"entry\":{\"databricks_cli_log\":{\"cli_test_event\":{\"name\":\"VALUE2\"}}}}",
			"{\"frontend_log_event_id\":\"a178892e-e285-4ce1-9114-55780875d64e\",\"entry\":{\"databricks_cli_log\":{\"cli_test_event\":{\"name\":\"VALUE3\"}}}}",
		}

		// Assert payload matches the expected payload.
		assert.Equal(m.t, expectedProtoLogs, reqb.ProtoLogs)
	}

	switch m.numCalls {
	case 1:
		// The call is successful but not all events are successfully logged.
		assertRequestPayload(request.(RequestBody))
		*(response.(*ResponseBody)) = ResponseBody{
			NumProtoSuccess: 3,
		}
	case 2:
		// The call is successful and all events are successfully logged.
		assertRequestPayload(request.(RequestBody))
		*(response.(*ResponseBody)) = ResponseBody{
			NumProtoSuccess: 4,
		}
	default:
		panic("unexpected number of calls")
	}

	return nil
}

func TestTelemetryLoggerFlushesEvents(t *testing.T) {
	mockClient := &mockDatabricksClient{
		t: t,
	}

	// Set the random number generator to a fixed seed to ensure that the UUIDs are deterministic.
	uuid.SetRand(rand.New(rand.NewSource(0)))
	t.Cleanup(func() {
		uuid.SetRand(nil)
	})

	ctx := WithDefaultLogger(context.Background())

	for _, v := range []events.DummyCliEnum{events.DummyCliEnumValue1, events.DummyCliEnumValue2, events.DummyCliEnumValue2, events.DummyCliEnumValue3} {
		Log(ctx, DatabricksCliLog{
			CliTestEvent: &events.CliTestEvent{Name: v},
		})
	}

	// Flush the events.
	Flush(ctx, mockClient)

	// Assert that the .Do method is called twice, because all logs were not
	// successfully logged in the first call.
	assert.Equal(t, 2, mockClient.numCalls)
}

func TestTelemetryLoggerFlushExitsOnTimeout(t *testing.T) {
	// Set the maximum additional wait time to 0 to ensure that the Flush method times out immediately.
	oldV := MaxAdditionalWaitTime
	MaxAdditionalWaitTime = 0
	t.Cleanup(func() {
		MaxAdditionalWaitTime = oldV
	})

	mockClient := &mockDatabricksClient{
		t: t,
	}

	// Set the random number generator to a fixed seed to ensure that the UUIDs are deterministic.
	uuid.SetRand(rand.New(rand.NewSource(0)))
	t.Cleanup(func() {
		uuid.SetRand(nil)
	})

	ctx := WithDefaultLogger(context.Background())

	for _, v := range []events.DummyCliEnum{events.DummyCliEnumValue1, events.DummyCliEnumValue2, events.DummyCliEnumValue2, events.DummyCliEnumValue3} {
		Log(ctx, DatabricksCliLog{
			CliTestEvent: &events.CliTestEvent{Name: v},
		})
	}

	// Flush the events.
	Flush(ctx, mockClient)

	// Assert that the .Do method is never called since the timeout is set to 0.
	assert.Equal(t, 0, mockClient.numCalls)
}