add tests for handleS3Operation

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()
+		})
+	}
+}

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, "<debug stack disabled>")
 	}
+	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)
+}