diff --git a/contrib/aws/internal/span_pointers/span_pointers_test.go b/contrib/aws/internal/span_pointers/span_pointers_test.go index 77ee4ce01d..5e95d36ad5 100644 --- a/contrib/aws/internal/span_pointers/span_pointers_test.go +++ b/contrib/aws/internal/span_pointers/span_pointers_test.go @@ -1,6 +1,17 @@ package span_pointers import ( + "context" + "encoding/json" + "github.com/aws/smithy-go/middleware" + smithyhttp "github.com/aws/smithy-go/transport/http" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gopkg.in/DataDog/dd-trace-go.v1/ddtrace" + "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer" + "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer" + "net/http" + "net/url" "testing" ) @@ -48,3 +59,127 @@ func TestGeneratePointerHash(t *testing.T) { }) } } + +func TestHandleS3Operation(t *testing.T) { + mt := mocktracer.Start() + defer mt.Stop() + + tests := []struct { + name string + bucket string + key string + etag string + expectedHash string + expectSuccess bool + }{ + { + name: "basic operation", + bucket: "some-bucket", + key: "some-key.data", + etag: "ab12ef34", + expectedHash: "e721375466d4116ab551213fdea08413", + expectSuccess: true, + }, + { + name: "quoted etag", + bucket: "some-bucket", + key: "some-key.data", + etag: "\"ab12ef34\"", + expectedHash: "e721375466d4116ab551213fdea08413", + expectSuccess: true, + }, + { + name: "non-ascii key", + bucket: "some-bucket", + key: "some-key.你好", + etag: "ab12ef34", + expectedHash: "d1333a04b9928ab462b5c6cadfa401f4", + expectSuccess: true, + }, + { + name: "empty bucket", + bucket: "", + key: "some_key", + etag: "some_etag", + expectSuccess: false, + }, + { + name: "empty key", + bucket: "some_bucket", + key: "", + etag: "some_etag", + expectSuccess: false, + }, + { + name: "empty etag", + bucket: "some_bucket", + key: "some_key", + etag: "", + expectSuccess: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := context.Background() + span, ctx := tracer.StartSpanFromContext(ctx, "test.s3.operation") + + // Create request + reqURL, _ := url.Parse("https://" + tt.bucket + ".s3.region.amazonaws.com/" + tt.key) + req := &smithyhttp.Request{ + Request: &http.Request{ + URL: reqURL, + }, + } + + // Create response + header := http.Header{} + header.Set("ETag", tt.etag) + res := &smithyhttp.Response{ + Response: &http.Response{ + Header: header, + }, + } + + // Create input/output + in := middleware.DeserializeInput{ + Request: req, + } + out := middleware.DeserializeOutput{ + RawResponse: res, + } + + AddSpanPointers("S3", in, out, span) + span.Finish() + spans := mt.FinishedSpans() + if tt.expectSuccess { + require.Len(t, spans, 1) + meta := spans[0].Tags() + + spanLinks, exists := meta["_dd.span_links"] + assert.True(t, exists, "Expected span links to be set") + assert.NotEmpty(t, spanLinks, "Expected span links to not be empty") + + spanLinksStr, ok := spanLinks.(string) + assert.True(t, ok, "Expected span links to be a string") + + var links []ddtrace.SpanLink + err := json.Unmarshal([]byte(spanLinksStr), &links) + require.NoError(t, err) + require.Len(t, links, 1) + + attributes := links[0].Attributes + assert.Equal(t, S3PointerKind, attributes["ptr.kind"]) + assert.Equal(t, PointerDownDirection, attributes["ptr.dir"]) + assert.Equal(t, LinkKind, attributes["link.kind"]) + assert.Equal(t, tt.expectedHash, attributes["ptr.hash"]) + } else { + require.Len(t, spans, 1) + tags := spans[0].Tags() + _, exists := tags["_dd.span_links"] + assert.False(t, exists, "Expected no span links to be set") + } + mt.Reset() + }) + } +} diff --git a/ddtrace/mocktracer/mockspan.go b/ddtrace/mocktracer/mockspan.go index 4701d96c7f..5a2578a5de 100644 --- a/ddtrace/mocktracer/mockspan.go +++ b/ddtrace/mocktracer/mockspan.go @@ -6,6 +6,7 @@ package mocktracer // import "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer" import ( + "encoding/json" "fmt" "sync" "time" @@ -206,6 +207,7 @@ func (s *mockspan) Finish(opts ...ddtrace.FinishOption) { if cfg.NoDebugStack { s.SetTag(ext.ErrorStack, "") } + s.serializeSpanLinksInMeta() s.Lock() defer s.Unlock() if s.finished { @@ -284,3 +286,18 @@ func (s *mockspan) Root() tracer.Span { root, _ := current.(*mockspan) return root } + +func (s *mockspan) AddSpanLink(link ddtrace.SpanLink) { + s.links = append(s.links, link) +} + +func (s *mockspan) serializeSpanLinksInMeta() { + if len(s.links) == 0 { + return + } + spanLinkBytes, err := json.Marshal(s.links) + if err != nil { + return + } + s.tags["_dd.span_links"] = string(spanLinkBytes) +}