From 93a5e9c1bad78078ed51e02f595df60fd86e1645 Mon Sep 17 00:00:00 2001
From: Deepak <dmallubhotla+github@gmail.com>
Date: Mon, 25 Jan 2021 14:50:54 -0600
Subject: [PATCH] Adds middleware set userid on context method.

---
 tokens/middleware.go              |  5 +++++
 tokens/middleware_context_test.go | 10 ++++++++++
 2 files changed, 15 insertions(+)

diff --git a/tokens/middleware.go b/tokens/middleware.go
index c8d9fa9..46b971e 100644
--- a/tokens/middleware.go
+++ b/tokens/middleware.go
@@ -67,6 +67,11 @@ func GetUserID(ctx context.Context) (int, error) {
 	return int(userID), nil
 }
 
+// SetUserID sets the username field on a context, necessary because the key is an unexported custom type.
+func SetUserID(ctx context.Context, id int) context.Context {
+	return context.WithValue(ctx, userIDCtxKey, int64(id))
+}
+
 // GetUsername does something similar to GetUserID.
 func GetUsername(ctx context.Context) (string, error) {
 	username, ok := ctx.Value(usernameCtxKey).(string)
diff --git a/tokens/middleware_context_test.go b/tokens/middleware_context_test.go
index f160280..06e2847 100644
--- a/tokens/middleware_context_test.go
+++ b/tokens/middleware_context_test.go
@@ -37,3 +37,13 @@ func TestBadContext(t *testing.T) {
 	assert.NotNil(err)
 
 }
+
+func TestSetContext(t *testing.T) {
+	assert := assert.New(t)
+
+	idToUse := 3
+	ctx := tokens.SetUserID(context.Background(), 3)
+	receivedID, err := tokens.GetUserID(ctx)
+	assert.Nil(err)
+	assert.EqualValues(idToUse, receivedID)
+}