Unit Testing with Context in Go

Often, I've needed to unit test a function that receives a context.Context as an argument. I've found that there are a few pitfalls to avoid when testing these functions, so I've written this up (maybe mainly for my future self) as a quick catalog of strategies for dealing with this sort of test. Below, I've built two related examples, one in which the function I'm testing does not directly interact with the passed in context and one in which the function does directly interact with the context.

Case #1: A simple pass-through

So, a lot of times with a REST request I'll have a context generated very early on in processing the request that I'll then pass on to lower and lower functions for logging, metrics, timeout, etc. There may be a number of functions along the way where I don't really directly need the context to do anything, but I want to pass it down to lower functions where I will need to access the context directly.

Here's an example of one such function:

func makeBackendRequest(ctx context.Context) (*http.Response, error) {
	req, _ := http.NewRequest(http.MethodGet, "http://backend-server.net", nil)
	req = req.WithContext(ctx)
	return client.Do(req)
}

Some notes:

• I've tried to use the standard library as much as possible to lower the amount of code needed. I'm ignoring the error returned by http.NewRequest for the same reason.
• This code just appends the passed in context onto a standard request and returns the response.
• The client here would be using the context for a timeout or deadline.

I'm a big fan of the auto unit test generation in Go, so I usually use that to start and then tailor the boilerplate to fit what I actually need. (In VS Code, assuming you've got the Go plugin installed, just right-click the function name and select Go: Generate Unit Tests for Function to do this). Here's what the auto-generated unit test looks like:

func Test_makeBackendRequest(t *testing.T) {
	type args struct {
		ctx context.Context
	}
	tests := []struct {
		name    string
		args    args
		want    *http.Response
		wantErr bool
	}{
		// TODO: Add test cases.
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := makeBackendRequest(tt.args.ctx)
			if (err != nil) != tt.wantErr {
				t.Errorf("makeBackendRequest() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("makeBackendRequest() = %v, want %v", got, tt.want)
			}
		})
	}
}

I don't really like the args struct here unless I have three or more arguments to pass into my function, so I'll get rid of that here. My normal move would be to just replace the args line in the test struct with the properties from the args struct definition at the top of the function, so that the start of the unit test function would look like:

func Test_makeBackendRequest(t *testing.T) {
	tests := []struct {
		name    string
		ctx     context.Context // used to be args
		want    *http.Response
		wantErr bool
	}{
		// TODO: Add test cases.
	}

However, since we are just passing context through here, it's value has no impact on any unit testing I'm doing on this function, so I can do even better. This setup above would lead me to create a variable like baseContext := context.Background() and then pass that in to each test case I create in my tests table. This would work just fine, but I do think that having that in the test table makes the test harder to read. In my opinion, only things which will change the output of the function (or describe the effect of the function like want and wantErr here) should be included in the test table. That way, it's much easier for me to use the test code to understand the intent of the functional code.¹

What we can do instead of declaring a baseContext variable is create that base context inside of the executing part of the unit test, the for loop which ranges over tests. We end up with something like:

func Test_makeBackendRequest(t *testing.T) {
	tests := []struct {
		name    string
		want    *http.Response
		wantErr bool
	}{
		// TODO: Add test cases.
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := context.Background()
			got, err := makeBackendRequest(ctx)
			if (err != nil) != tt.wantErr {
				t.Errorf("makeBackendRequest() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("makeBackendRequest() = %v, want %v", got, tt.want)
			}
		})
	}
}

So much for the simplest case.

Case #2: Need to actually do something with context in test

This can be easily extended, then, if you need to test something about the context. Imagine that this is what the function we are unit testing looks like:

func makeBackendRequest(ctx context.Context) (*http.Response, error) {
	rawMethod := ctx.Value("method")
	reqMethod, ok := rawMethod.(string)
	if !ok {
		return nil, errors.New("invalid method on context")
	}

	req, _ := http.NewRequest(reqMethod, "http://backend-server.net", nil)
	req = req.WithContext(ctx)
	return client.Do(req)
}

Now, for some reason, we're pulling the request method for our backend request out of context,² and we need to write some test cases: one where the context contains a nil value for "method", one where it contains a non-string value for "method", and one where it contains a valid string value for "method".³ Since the only thing that's changing here is the value in context, we can still use the same mechanism for creating context in each test run. We'll just need to add a step where we populate the value with one we've prepared in our test cases.

So first we'll add the value, which I'll call method, to our test struct using an empty interface{} here so that we can put both string and non-string values in.

tests := []struct {
		name    string
		method  interface{}
		want    *http.Response
		wantErr bool
	}{
		// TODO: Add test cases.
	}

Now we'll have our test execution section store whatever value is in method for each test into the context we generate:

for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := context.Background()
			ctx = context.WithValue(ctx, "method", tt.method)
			got, err := makeBackendRequest(ctx)
			[ ... ]
		})
	}

And now we can finally create the test cases we laid out above. So the first case should reflect a nil or missing value in the context:

tests := []struct {
		name   string
		method interface{}
		// want    *http.Response
		wantErr bool
	}{
		{
			name:    "value is nil",
			method:  nil,
			wantErr: true,
		},
	}

I've commented out the want line in my test struct because I don't really know what I want in this example and setting up dummy *http.Responses would be good for a post of its own. I also want to point out that I am explicitly setting the value of method in my test case to nil. I could implicitly set this value by leaving the assignment of a value to method out, because the zero (default) value for an interface{} is nil. I've worked on teams with people who prefer that because it is less typing, but I'm a firm believer in being explicit here. One of my favorite things about Go is that, as Peter Bourgon describes it, “it is basically non-magical”, so, whenever there's a style choice to be made, I try to choose whichever style most matches this characteristic of Go. Implicit-ness tends toward magic, in my opinion, so I'm explicit here. This is certainly a matter of personal or team preference, though, as either way of setting method to nil here has the same technical effect.

All that remains is to add the other two test cases. Here's the final full test code:

(Note: I've commented out all of the code related to checking the primary return value)

func Test_makeBackendRequest(t *testing.T) {
	tests := []struct {
		name   string
		method interface{}
		// want    *http.Response
		wantErr bool
	}{
		{
			name:    "value is nil",
			method:  nil,
			wantErr: true,
		},
		{
			name:    "value is invalid - integer",
			method:  1,
			wantErr: true,
		},
		{
			name:    "value is valid string",
			method:  http.MethodGet,
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := context.Background()
			ctx = context.WithValue(ctx, "method", tt.method)
			_, err := makeBackendRequest(ctx)
			if (err != nil) != tt.wantErr {
				t.Errorf("makeBackendRequest() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			// if !reflect.DeepEqual(got, tt.want) {
			// 	t.Errorf("makeBackendRequest() = %v, want %v", got, tt.want)
			// }
		})
	}
}

This same approach would extend out if you had code dealing with cancels, deadlines, or timeouts; simply set up values for your CancelFunc and any deadlines or timeouts needed in your test cases, then when you create your context in the functional bit, use WithCancel, WithDeadline, or WithTimeout as we did WithValue.