completed unittests for users package

example-project/internal/user/auth_test.go

@@ -15,31 +15,33 @@ import (
 
 // mockTokenGenerator is used for testing that Authenticate calls its provided
 // token generator in a specific way.
-type mockTokenGenerator struct{}
-
-// Private key generated by GenerateToken that is need for ParseClaims
-var mockTokenKey *rsa.PrivateKey
+type mockTokenGenerator struct {
+	// Private key generated by GenerateToken that is need for ParseClaims
+	key *rsa.PrivateKey
+	// algorithm is the method used to generate the private key.
+	algorithm string
+}
 
 // GenerateToken implements the TokenGenerator interface. It returns a "token"
 // that includes some information about the claims it was passed.
-func (g mockTokenGenerator) GenerateToken(claims auth.Claims) (string, error) {
+func (g *mockTokenGenerator) GenerateToken(claims auth.Claims) (string, error) {
 	privateKey, err := auth.Keygen()
 	if err != nil {
 		return "", err
 	}
 
-	mockTokenKey, err = jwt.ParseRSAPrivateKeyFromPEM(privateKey)
+	g.key, err = jwt.ParseRSAPrivateKeyFromPEM(privateKey)
 	if err != nil {
 		return "", err
 	}
 
-	algorithm := "RS256"
-	method := jwt.GetSigningMethod(algorithm)
+	g.algorithm = "RS256"
+	method := jwt.GetSigningMethod(g.algorithm)
 
 	tkn := jwt.NewWithClaims(method, claims)
 	tkn.Header["kid"] = "1"
 
-	str, err := tkn.SignedString(mockTokenKey)
+	str, err := tkn.SignedString(g.key)
 	if err != nil {
 		return "", err
 	}
@@ -49,18 +51,17 @@ func (g mockTokenGenerator) GenerateToken(claims auth.Claims) (string, error) {
 
 // ParseClaims recreates the Claims that were used to generate a token. It
 // verifies that the token was signed using our key.
-func (g mockTokenGenerator) ParseClaims(tknStr string) (auth.Claims, error) {
-	algorithm := "RS256"
+func (g *mockTokenGenerator) ParseClaims(tknStr string) (auth.Claims, error) {
 	parser := jwt.Parser{
-		ValidMethods: []string{algorithm},
+		ValidMethods: []string{g.algorithm},
 	}
 
-	if mockTokenKey == nil {
-		panic("key is nil")
+	if g.key == nil {
+		return auth.Claims{}, errors.New("Private key is empty.")
 	}
 
 	f := func(t *jwt.Token) (interface{}, error) {
-		return mockTokenKey.Public().(*rsa.PublicKey), nil
+		return g.key.Public().(*rsa.PublicKey), nil
 	}
 
 	var claims auth.Claims
@@ -93,7 +94,7 @@ func TestAuthenticate(t *testing.T) {
 			now := time.Now().Add(time.Hour * -1)
 
 			// Try to authenticate an invalid user.
-			_, err := Authenticate(ctx, test.MasterDB, tknGen, now, "doesnotexist@gmail.com", "xy7")
+			_, err := Authenticate(ctx, test.MasterDB, tknGen, "doesnotexist@gmail.com", "xy7", time.Hour, now)
 			if errors.Cause(err) != ErrAuthenticationFailure {
 				t.Logf("\t\tGot : %+v", err)
 				t.Logf("\t\tWant: %+v", ErrAuthenticationFailure)
@@ -117,23 +118,36 @@ func TestAuthenticate(t *testing.T) {
 
 			// Create a new random account and associate that with the user.
 			// This defined role should be the claims.
-			accountId := uuid.NewRandom().String()
-			accountRole := UserAccountRole_Admin
+			account1Id := uuid.NewRandom().String()
+			account1Role := UserAccountRole_Admin
 			_, err = AddAccount(tests.Context(), auth.Claims{}, test.MasterDB, AddAccountRequest{
 				UserID:    user.ID,
-				AccountID: accountId,
-				Roles:     []UserAccountRole{accountRole},
+				AccountID: account1Id,
+				Roles:     []UserAccountRole{account1Role},
 			}, now)
 			if err != nil {
 				t.Log("\t\tGot :", err)
 				t.Fatalf("\t%s\tAddAccount failed.", tests.Failed)
 			}
 
+			// Create a second new random account and associate that with the user.
+			account2Id := uuid.NewRandom().String()
+			account2Role := UserAccountRole_User
+			_, err = AddAccount(tests.Context(), auth.Claims{}, test.MasterDB, AddAccountRequest{
+				UserID:    user.ID,
+				AccountID: account2Id,
+				Roles:     []UserAccountRole{account2Role},
+			}, now.Add(time.Second))
+			if err != nil {
+				t.Log("\t\tGot :", err)
+				t.Fatalf("\t%s\tAddAccount failed.", tests.Failed)
+			}
+
 			// Add 30 minutes to now to simulate time passing.
 			now = now.Add(time.Minute * 30)
 
 			// Try to authenticate valid user with invalid password.
-			_, err = Authenticate(ctx, test.MasterDB, tknGen, now, user.Email, "xy7")
+			_, err = Authenticate(ctx, test.MasterDB, tknGen, user.Email, "xy7", time.Hour, now)
 			if errors.Cause(err) != ErrAuthenticationFailure {
 				t.Logf("\t\tGot : %+v", err)
 				t.Logf("\t\tWant: %+v", ErrAuthenticationFailure)
@@ -142,7 +156,7 @@ func TestAuthenticate(t *testing.T) {
 			t.Logf("\t%s\tAuthenticate user w/invalid password ok.", tests.Success)
 
 			// Verify that the user can be authenticated with the created user.
-			tkn, err := Authenticate(ctx, test.MasterDB, tknGen, now, user.Email, initPass)
+			tkn1, err := Authenticate(ctx, test.MasterDB, tknGen, user.Email, initPass, time.Hour, now)
 			if err != nil {
 				t.Log("\t\tGot :", err)
 				t.Fatalf("\t%s\tAuthenticate user failed.", tests.Failed)
@@ -150,18 +164,40 @@ func TestAuthenticate(t *testing.T) {
 			t.Logf("\t%s\tAuthenticate user ok.", tests.Success)
 
 			// Ensure the token string was correctly generated.
-			claims, err := tknGen.ParseClaims(tkn.Token)
+			claims1, err := tknGen.ParseClaims(tkn1.Token)
 			if err != nil {
 				t.Log("\t\tGot :", err)
 				t.Fatalf("\t%s\tParse claims from token failed.", tests.Failed)
-			} else if diff := cmp.Diff(claims, tkn.claims); diff != "" {
+			} else if diff := cmp.Diff(claims1, tkn1.claims); diff != "" {
 				t.Fatalf("\t%s\tExpected parsed claims to match from token. Diff:\n%s", tests.Failed, diff)
-			} else if diff := cmp.Diff(claims.Roles, []string{accountRole.String()}); diff != "" {
+			} else if diff := cmp.Diff(claims1.Roles, []string{account1Role.String()}); diff != "" {
 				t.Fatalf("\t%s\tExpected parsed claims roles to match user account. Diff:\n%s", tests.Failed, diff)
-			} else if diff := cmp.Diff(claims.AccountIds, []string{accountId}); diff != "" {
+			} else if diff := cmp.Diff(claims1.AccountIds, []string{account1Id, account2Id}); diff != "" {
 				t.Fatalf("\t%s\tExpected parsed claims account IDs to match the single user account. Diff:\n%s", tests.Failed, diff)
 			}
-			t.Logf("\t%s\tParse claims from token ok.", tests.Success)
+			t.Logf("\t%s\tAuthenticate parse claims from token ok.", tests.Success)
+
+			// Try switching to a second account using the first set of claims.
+			tkn2, err := SwitchAccount(ctx, test.MasterDB, tknGen, claims1, account2Id, time.Hour, now)
+			if err != nil {
+				t.Log("\t\tGot :", err)
+				t.Fatalf("\t%s\tSwitchAccount user failed.", tests.Failed)
+			}
+			t.Logf("\t%s\tSwitchAccount user ok.", tests.Success)
+
+			// Ensure the token string was correctly generated.
+			claims2, err := tknGen.ParseClaims(tkn2.Token)
+			if err != nil {
+				t.Log("\t\tGot :", err)
+				t.Fatalf("\t%s\tParse claims from token failed.", tests.Failed)
+			} else if diff := cmp.Diff(claims2, tkn2.claims); diff != "" {
+				t.Fatalf("\t%s\tExpected parsed claims to match from token. Diff:\n%s", tests.Failed, diff)
+			} else if diff := cmp.Diff(claims2.Roles, []string{account2Role.String()}); diff != "" {
+				t.Fatalf("\t%s\tExpected parsed claims roles to match user account. Diff:\n%s", tests.Failed, diff)
+			} else if diff := cmp.Diff(claims2.AccountIds, []string{account1Id, account2Id}); diff != "" {
+				t.Fatalf("\t%s\tExpected parsed claims account IDs to match the single user account. Diff:\n%s", tests.Failed, diff)
+			}
+			t.Logf("\t%s\tSwitchAccount parse claims from token ok.", tests.Success)
 		}
 	}
 }