From ce17497f3c9f94a7cd6cccefbc570e2ed1870720 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 30 Oct 2024 15:52:47 +0100 Subject: [PATCH] add unit test --- cmd/root/user_agent_upstream_test.go | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/cmd/root/user_agent_upstream_test.go b/cmd/root/user_agent_upstream_test.go index fc6ea0c75..acc226038 100644 --- a/cmd/root/user_agent_upstream_test.go +++ b/cmd/root/user_agent_upstream_test.go @@ -2,10 +2,14 @@ package root import ( "context" + "regexp" "testing" "github.com/databricks/databricks-sdk-go/useragent" + "github.com/google/uuid" + "github.com/spf13/cobra" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestUpstreamSet(t *testing.T) { @@ -43,3 +47,20 @@ func TestUpstreamVersionSetUpstreamNotSet(t *testing.T) { assert.NotContains(t, useragent.FromContext(ctx), "upstream/") assert.NotContains(t, useragent.FromContext(ctx), "upstream-version/") } + +func TestWithCommandInUserAgent(t *testing.T) { + ctx := withCommandInUserAgent(context.Background(), &cobra.Command{Use: "foo"}) + + // Check that the command name is in the user agent string. + ua := useragent.FromContext(ctx) + assert.Contains(t, ua, "cmd/foo") + + // Check that the command trace ID is in the user agent string. + re := regexp.MustCompile(`command-trace-id/([a-f0-9-]+) `) + matches := re.FindAllStringSubmatch(ua, -1) + + // Assert that we have exactly one match and that it's a valid UUID. + require.Len(t, matches, 1) + _, err := uuid.Parse(matches[0][1]) + assert.NoError(t, err) +}