Skip to content

Commit

Permalink
adding tests and refactoring
Browse files Browse the repository at this point in the history
  • Loading branch information
@puneet-traceable
puneet-traceable committed Nov 11, 2023
1 parent 6f98475 commit e3beab5
Show file tree
Hide file tree
Showing 2 changed files with 168 additions and 40 deletions.
60 changes: 20 additions & 40 deletions sdk/instrumentation/bodyattribute/bodyattribute.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,75 +17,55 @@ func SetTruncatedBodyAttribute(attrName string, body []byte, bodyMaxSize int, sp
}

if bodyLen <= bodyMaxSize {
span.SetAttribute(attrName, string(body))
SetBodyAttribute(attrName, body, false, span)
return
}

truncatedBody := body[:bodyMaxSize]
span.SetAttribute(fmt.Sprintf("%s.truncated", attrName), true)
span.SetAttribute(attrName, string(truncatedBody))
SetBodyAttribute(attrName, body[:bodyMaxSize], true, span)
}

// SetTruncatedEncodedBodyAttribute is like SetTruncatedBodyAttribute above but also base64 encodes the
// body. This is usually due to non utf8 bytes in the body eg. for multipart/form-data content type.
// The body attribute name has a ".base64" suffix.
func SetTruncatedEncodedBodyAttribute(attrName string, body []byte, bodyMaxSize int, span sdk.Span) {
bodyLen := len(body)
if bodyLen == 0 {
if len(body) == 0 {
return
}

if bodyLen <= bodyMaxSize {
span.SetAttribute(attrName+".base64", base64.RawStdEncoding.EncodeToString(body))
SetEncodedBodyAttribute(attrName, body, false, span)
return
}

truncatedBody := body[:bodyMaxSize]
span.SetAttribute(fmt.Sprintf("%s.truncated", attrName), true)
span.SetAttribute(attrName+".base64", base64.RawStdEncoding.EncodeToString(truncatedBody))
SetEncodedBodyAttribute(attrName, body[:bodyMaxSize], true, span)
}

// SetPossiblyTruncatedBodyAttribute sets the body as a span attribute.
// When body is being truncated, or already truncated we also add a second attribute suffixed by `.truncated` to
// make it clear to the user, body has been modified.
func SetPossiblyTruncatedBodyAttribute(attrName string, body []byte, bodyMaxSize int, truncated bool, span sdk.Span) {
bodyLen := len(body)
if bodyLen == 0 {
// SetBodyAttribute sets the body as a span attribute.
// also sets truncated attribute if truncated is true
func SetBodyAttribute(attrName string, body []byte, truncated bool, span sdk.Span) {
if len(body) == 0 {
return
}

if bodyLen <= bodyMaxSize {
span.SetAttribute(attrName, string(body))
// if already truncated then set attribute
if truncated {
span.SetAttribute(fmt.Sprintf("%s.truncated", attrName), true)
}
return
span.SetAttribute(attrName, string(body))
// if already truncated then set attribute
if truncated {
span.SetAttribute(fmt.Sprintf("%s.truncated", attrName), true)
}

truncatedBody := body[:bodyMaxSize]
span.SetAttribute(fmt.Sprintf("%s.truncated", attrName), true)
span.SetAttribute(attrName, string(truncatedBody))
}

// SetPossiblyTruncatedEncodedBodyAttribute is like SetPossiblyTruncatedBodyAttribute above but also base64 encodes the
// SetEncodedBodyAttribute is like SetBodyAttribute above but also base64 encodes the
// body. This is usually due to non utf8 bytes in the body eg. for multipart/form-data content type.
// The body attribute name has a ".base64" suffix.
func SetPossiblyTruncatedEncodedBodyAttribute(attrName string, body []byte, bodyMaxSize int, truncated bool, span sdk.Span) {
bodyLen := len(body)
if bodyLen == 0 {
func SetEncodedBodyAttribute(attrName string, body []byte, truncated bool, span sdk.Span) {
if len(body) == 0 {
return
}

if bodyLen <= bodyMaxSize {
span.SetAttribute(attrName+".base64", base64.RawStdEncoding.EncodeToString(body))
if truncated {
span.SetAttribute(fmt.Sprintf("%s.truncated", attrName), true)
}
return
span.SetAttribute(attrName+".base64", base64.RawStdEncoding.EncodeToString(body))
// if already truncated then set attribute
if truncated {
span.SetAttribute(fmt.Sprintf("%s.truncated", attrName), true)
}

truncatedBody := body[:bodyMaxSize]
span.SetAttribute(fmt.Sprintf("%s.truncated", attrName), true)
span.SetAttribute(attrName+".base64", base64.RawStdEncoding.EncodeToString(truncatedBody))
}
148 changes: 148 additions & 0 deletions sdk/instrumentation/bodyattribute/bodyattribute_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,3 +51,151 @@ func TestSetTruncatedEncodedBodyAttributeEmptyBody(t *testing.T) {
assert.Nil(t, s.ReadAttribute("http.request.body.base64"))
assert.Zero(t, s.RemainingAttributes())
}

func TestSetBodyAttribute(t *testing.T) {
testBody := "test1test2"
type args struct {
attrName string
body []byte
truncated bool
span *mock.Span
}
tests := []struct {
name string
args args
expectedAssertions func(t *testing.T, gotSpan *mock.Span)
}{
{
name: "empty body, truncated",
args: args{
attrName: "http.request.body",
body: []byte(""),
truncated: true,
span: mock.NewSpan(),
},
expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
assert.Nil(t, gotSpan.ReadAttribute("http.request.body"))
assert.Zero(t, gotSpan.RemainingAttributes())
},
},
{
name: "empty body, not truncated",
args: args{
attrName: "http.request.body",
body: []byte(""),
truncated: false,
span: mock.NewSpan(),
},
expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
assert.Nil(t, gotSpan.ReadAttribute("http.request.body"))
assert.Zero(t, gotSpan.RemainingAttributes())
},
},
{
name: "non empty body, not truncated",
args: args{
attrName: "http.request.body",
body: []byte(testBody),
truncated: false,
span: mock.NewSpan(),
},
expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
assert.Equal(t, testBody, gotSpan.ReadAttribute("http.request.body"))
assert.Zero(t, gotSpan.RemainingAttributes())
},
},
{
name: "non empty body, truncated",
args: args{
attrName: "http.request.body",
body: []byte(testBody),
truncated: true,
span: mock.NewSpan(),
},
expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
assert.Equal(t, testBody, gotSpan.ReadAttribute("http.request.body"))
assert.True(t, (gotSpan.ReadAttribute("http.request.body.truncated")).(bool))
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
SetBodyAttribute(tt.args.attrName, tt.args.body, tt.args.truncated, tt.args.span)
tt.expectedAssertions(t, tt.args.span)
})
}
}

func TestSetEncodedBodyAttribute(t *testing.T) {
testBody := "test1test2"
type args struct {
attrName string
body []byte
truncated bool
span *mock.Span
}
tests := []struct {
name string
args args
expectedAssertions func(t *testing.T, gotSpan *mock.Span)
}{
{
name: "empty body, truncated",
args: args{
attrName: "http.request.body",
body: []byte(""),
truncated: true,
span: mock.NewSpan(),
},
expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
assert.Nil(t, gotSpan.ReadAttribute("http.request.body.base64"))
assert.Zero(t, gotSpan.RemainingAttributes())
},
},
{
name: "empty body, not truncated",
args: args{
attrName: "http.request.body",
body: []byte(""),
truncated: false,
span: mock.NewSpan(),
},
expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
assert.Nil(t, gotSpan.ReadAttribute("http.request.body.base64"))
assert.Zero(t, gotSpan.RemainingAttributes())
},
},
{
name: "non empty body, not truncated",
args: args{
attrName: "http.request.body",
body: []byte(testBody),
truncated: false,
span: mock.NewSpan(),
},
expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
assert.Equal(t, base64.RawStdEncoding.EncodeToString([]byte(testBody)), gotSpan.ReadAttribute("http.request.body.base64"))
assert.Zero(t, gotSpan.RemainingAttributes())
},
},
{
name: "non empty body, truncated",
args: args{
attrName: "http.request.body",
body: []byte(testBody),
truncated: true,
span: mock.NewSpan(),
},
expectedAssertions: func(t *testing.T, gotSpan *mock.Span) {
assert.Equal(t, base64.RawStdEncoding.EncodeToString([]byte(testBody)), gotSpan.ReadAttribute("http.request.body.base64"))
assert.True(t, (gotSpan.ReadAttribute("http.request.body.truncated")).(bool))
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
SetEncodedBodyAttribute(tt.args.attrName, tt.args.body, tt.args.truncated, tt.args.span)
tt.expectedAssertions(t, tt.args.span)
})
}
}

0 comments on commit e3beab5

Please sign in to comment.