Unit Testing Annotations

[mousedownmike]

In my unit tests, I'm trying to verify x-ray annotations are added to a Segment. I'm having difficulty finding a reliable way to achieve this. The code I have in this issue can also be found in this minimal repository.

The scenario this illustrates is:

• A handler function is called by some component that has already started a Segment (think Lambda)
• The function creates a named SubSegment
• The function uses the context to do "stuff" (S3, Dynamo, etc...)
• The function adds an Annotation to the Segment.
• The SubSegment is closed

If I run a single test (e.g. TestMyHandler_Get) by itself, it will pass. If I run all the tests in the file, they will start failing at the assertions for the annotations. I have tried a number of things to get this verified but none has worked so far:

• I have put the annotation in the SubSegment (instead of the ParentSegment), closed the subsegment and then closed the parent segment in the unit test. From there I have used the Subsegments : []json.RawMessage value from the segment to check for the existence of the annotation parts (i.e. key and value separately). Again, if I run one test, it works but if I run them all it fails.
• I have tried various combinations of using context.Background() or the context value returned from BeginSubsegment() thinking there was a Context problem but nothing changed.
• I have verified that the data IS successfully submitted via the Daemon to the X-Ray console and annotations appear on either the SubSegment or Segment depending on how I set them.

For now, I'm pulling this verification out but it would be immensely helpful to ensure consistent annotations across my various handlers. Is there a practical way of testing if a Segment and/or SubSegment contains annotations added during a unit test?

Any help is greatly appreciated.

pkg/api/handler.go

type MyHandler struct {
    succeed bool
}

func (h MyHandler) Get(w http.ResponseWriter, r *http.Request) {
    ctx, xs := xray.BeginSubsegment(r.Context(), "MyHandler.Get")
    h.useContext(ctx)
    _ = xs.ParentSegment.AddAnnotation("method", "GET")
    if h.succeed {
        w.WriteHeader(http.StatusOK)
    } else {
        w.WriteHeader(http.StatusInternalServerError)
    }
    _, e := w.Write([]byte("GET MyHandler"))
    xs.Close(e)
}

func (h MyHandler) Post(w http.ResponseWriter, r *http.Request) {
    ctx, xs := xray.BeginSubsegment(r.Context(), "MyHandler.Post")
    h.useContext(ctx)
    _ = xs.ParentSegment.AddAnnotation("method", "POST")
    if h.succeed {
        w.WriteHeader(http.StatusOK)
    } else {
        w.WriteHeader(http.StatusInternalServerError)
    }
    _, e := w.Write([]byte("POST MyHandler"))
    xs.Close(e)
}

func (h MyHandler) useContext(_ context.Context) {
    // make calls somewhere with the context
}

pkg/api/handler_test.go

func TestMyHandler_Get(t *testing.T) {
    var h MyHandler
    var w *httptest.ResponseRecorder
    var r *http.Request

    mockCtx, xs := xray.BeginSegment(context.Background(), t.Name())
    defer xs.Close(nil)
    before := func() {
        h = MyHandler{}
        w = httptest.NewRecorder()
        r = &http.Request{}
        r = r.WithContext(mockCtx)
    }

    t.Run("success", func(t *testing.T) {
        before()
        h.succeed = true
        h.Get(w, r)
        assert.Equal(t, http.StatusOK, w.Code)
        assert.Equal(t, "GET", xs.Annotations["method"])
    })
    t.Run("fail", func(t *testing.T) {
        before()
        h.succeed = false
        h.Get(w, r)
        assert.Equal(t, http.StatusInternalServerError, w.Code)
        assert.Equal(t, "GET", xs.Annotations["method"])
    })
}

func TestMyHandler_Post(t *testing.T) {
    var h MyHandler
    var w *httptest.ResponseRecorder
    var r *http.Request

    mockCtx, xs := xray.BeginSegment(context.Background(), t.Name())
    defer xs.Close(nil)
    before := func() {
        h = MyHandler{}
        w = httptest.NewRecorder()
        r = &http.Request{}
        r = r.WithContext(mockCtx)
    }

    t.Run("success", func(t *testing.T) {
        before()
        h.succeed = true
        h.Post(w, r)
        assert.Equal(t, http.StatusOK, w.Code)
        assert.Equal(t, "POST", xs.Annotations["method"])
    })
    t.Run("fail", func(t *testing.T) {
        before()
        h.succeed = false
        h.Post(w, r)
        assert.Equal(t, http.StatusInternalServerError, w.Code)
        assert.Equal(t, "POST", xs.Annotations["method"])
    })
}

[bhautikpip]

Hi @mousedownmike ,

I really appreciate you providing detailed explanation. I was able to reproduce this issue and also was able to find the cause for this. So Basically, this is a sampling issue. When you are running your testing suite together it runs TestMyHandler_Get and it runs fine but it always fails in TestMyHandler_Post because at that time sampling value is false and because of this logic in SDK (https://github.com/aws/aws-xray-sdk-go/blob/95821db08ed031914e88b3a7f06c95a943679375/xray/segment.go#L505) it never adds annotation when sampling is false so that's why you're getting nil value when running TestMyHandler_Post.

There are couple workarounds.

1. add 1 second delay in your second test case (time.Sleep(1 * time.Second))
2. Make it 100% sampling (while running this testing suite)

[mousedownmike]

@bhautikpip ,

Thank you for the reply and sorry it has taken me so long to get back to this. Following your direction, I ended up creating:

type AlwaysSamplingStrategy struct {}
func (s *AlwaysSamplingStrategy) ShouldTrace(_ *sampling.Request) *sampling.Decision {
    return &sampling.Decision{Sample: true}
}

And I introduced it into my tests with:

ctx, _ = xray.ContextWithConfig(context.Background(), xray.Config{SamplingStrategy: &AlwaysSamplingStrategy{}})

When I do that I can successfully assert the values of annotations that are placed on the ParentSegment as I described in my original request. Ideally I would like to be able to leave the Annotations on the Subsegment (i.e. xs.AddAnnotation("method", "GET") instead of xs.ParentSegment.AddAnnotation("method", "GET")

At the time of my assertions, the value of xs.Subsegments is nil but I do see the values in xs.rawSubsegments. Unfortunately that isn't visible to me. Maybe I should open a separate ticket for this, but can you think of a way I can assert the values of the Subsegment annotations?

Thanks!

[bhautikpip]

Hi @mousedownmike ,

I think if you add annotations at subsegment level then you should be able to validate those as well. Would you be able to post your repro code and the error you're receiving when you do that?

[mousedownmike]

Thanks for checking in @bhautikpip

I guess my primary issue is that I don't understand how to access the Subsegments from the parent segment in my test code. I would think something like the following line would work:

assert.Equal(t, "SUB-GET", xs.Subsegments[0].Annotations["subsegment_annotation"])

But that isn't valid. Here are some updated code snippets that assume the AlwaysSamplingStrategy is in use and illustrate what works and where I get stuck.

pkg/api/handler.go (see earlier issue for more context)

func (h MyHandler) Get(w http.ResponseWriter, r *http.Request) {
    ctx, xs := xray.BeginSubsegment(r.Context(), "MyHandler.Get")
    _ = xs.ParentSegment.AddAnnotation("parent_segment_annotation", "GET")
    _ = xs.AddAnnotation("subsegment_annotation", "SUB-GET")
    if h.succeed {
        w.WriteHeader(http.StatusOK)
    } else {
        w.WriteHeader(http.StatusInternalServerError)
    }
    _, e := w.Write([]byte("GET MyHandler"))
    xs.Close(e)
}

pkg/api/handler_test.go (see earlier issue for more context)

func TestMyHandler_Get(t *testing.T) {
    var h MyHandler
    var w *httptest.ResponseRecorder
    var r *http.Request

    ctx, _ := xray.ContextWithConfig(context.Background(), xray.Config{SamplingStrategy: &AlwaysSamplingStrategy{}})
    ctx, xs := xray.BeginSegment(ctx, t.Name())
    defer xs.Close(nil)
    before := func() {
        h = MyHandler{}
        w = httptest.NewRecorder()
        r = &http.Request{}
        r = r.WithContext(ctx)
    }

    t.Run("success", func(t *testing.T) {
        before()
        h.succeed = true
        h.Get(w, r)
        assert.Equal(t, http.StatusOK, w.Code)
        assert.Equal(t, "GET", xs.Annotations["parent_segment_annotation"]) // Works!
        // The following is not valid code.  Subsegments is really a []json.RawMessage and also nil when this running
        assert.Equal(t, "SUB-GET", xs.Subsegments[0].Annotations["subsegment_annotation"])
    })
...

[bhautikpip]

Hi @mousedownmike ,

I was able to try out your issue. I think you're having difficulty accessing subsegment from segment level is because Subsegment value in segment's data structure will be populated once it runs segment.Close() method. So if you see the value before seg.Close() method it will be nil and you're right about rawSubsegments which will contain subsegments but unfortunately you won't be able to access that using segment. My follow up question would be is there any reason why you're not directly accessing sub segment's annotation value by accessing subsegment in your test ?

[mousedownmike]

Hi @bhautikpip

Regarding:

directly accessing sub segment's annotation value by accessing subsegment in your test

I think that's what I'm trying to do but I just don't know how to get a handle to the subsegment in my test. The subsegment is created (and closed) in the handler so I'm not sure how to get a handle to it without going through the segment I'm receiving from ctx, xs := xray.BeginSegment(ctx, t.Name()). I create the parent segment in the test to emulate the context that will be in place when the handler is run in a Lambda function.

[bhautikpip]

Hi @mousedownmike ,

I think you probably can use below example to get access to subsegment and then access the value of annotation key.

package main

import (
    "context"
    "fmt"
    "github.com/aws/aws-xray-sdk-go/xray"
)

func main() {
    ctx, seg := xray.BeginSegment(context.Background(), "seg")
    subSegCtx, subSeg := xray.BeginSubsegment(ctx, "subseg")
    _ = subSeg.AddAnnotation("key", "value")
    subSeg.Close(nil)
    sub := xray.GetSegment(subSegCtx)
    fmt.Println(sub)
    seg.Close(nil)
}

[mousedownmike]

Thanks @bhautikpip, unfortunately it doesn't look like this will work because the subSegCtx is only available inside my handler method and not exposed to the unit test. xray.GetSegment(ctx) does get me a segment that has rawSubsegments with the data I'm looking for but it isn't available either (only visible in the debugger).

The next best thing I can think of is possibly mocking the Emitter interface to see if I can record the subsegment there.

[bhautikpip]

Yeah, in order to use above example you need subsegment context. I think probably mocking Emitter interface should work if you try to access segment from there.

[mousedownmike]

@bhautikpip it doesn't look like segments get emitted in time for me to assert them. I'm starting to feel like this sort of testing may just be impossible with the current library. Unless you have other ideas, I'm inclined to just close this issue.

[bhautikpip]

Yeah, timing would be an issue in the case of emitter. Also, without context it would be hard to access subsegment's fields. I would agree that testing you're trying to achieve would be hard to do with the current state of library.

[mousedownmike]

Thanks for working with me on this. It was always in the "nice to have" bucket so I'm going to close this issue.