diff --git a/cmd/cognito_test/jwtstuff/main.go b/cmd/cognito_test/jwtstuff/main.go
index ec23883a..98781cd8 100644
--- a/cmd/cognito_test/jwtstuff/main.go
+++ b/cmd/cognito_test/jwtstuff/main.go
@@ -3,91 +3,107 @@ package main
 import (
 	"fmt"
 	"log"
-	"queryorchestration/internal/rbac" // Replace with your actual module name
+
+	"queryorchestration/internal/rbac"
 	"queryorchestration/internal/serviceconfig"
+
+	"github.com/golang-jwt/jwt/v5"
 )
 
-// for testing run like `LOG_LEVEL=INFO go run main.go`
-func main() {
-
+func setupConfig() (*serviceconfig.BaseConfig, error) {
 	cfg := &serviceconfig.BaseConfig{}
-	errGettingConfig := serviceconfig.InitializeConfig(cfg)
-	if errGettingConfig != nil {
-		log.Fatal(errGettingConfig)
+	if err := serviceconfig.InitializeConfig(cfg); err != nil {
+		return nil, fmt.Errorf("failed to initialize config: %w", err)
 	}
 
-	errorInitializingAuthProvider := cfg.InitializeAuthProvider()
-	if errorInitializingAuthProvider != nil {
-		log.Fatal(errorInitializingAuthProvider)
+	if err := cfg.InitializeAuthProvider(); err != nil {
+		return nil, fmt.Errorf("failed to initialize auth provider: %w", err)
 	}
 
 	cfg.PrintConfig("secret")
-	// use the config here for the key provider parameters
-	// Example with local key provider for testing
 	if cfg.AuthProvider == nil {
-		log.Fatal("AuthProvider is nil")
+		return nil, fmt.Errorf("auth provider is nil")
 	}
 
-	localProvider := rbac.NewLocalKeyProvider("private_key.pem", "public_key.pem")
+	return cfg, nil
+}
 
-	// Generate a mock JWT with multiple Cognito groups
+func generateAndValidateToken(localProvider *rbac.LocalKeyProvider) (*rbac.TokenDetails, error) {
 	groups := []string{"admins", "developers", "querybuilders", "uploaders", "exporters"}
 	tokenString, err := rbac.GenerateTestJWT(localProvider, "test-user-id", "test@example.com", groups)
 	if err != nil {
-		fmt.Println("Error generating token:", err)
-		return
+		return nil, fmt.Errorf("error generating token: %w", err)
 	}
 
 	fmt.Printf("raw token: %s\n", tokenString)
 
 	token, claims, err := rbac.ValidateJWT(tokenString, localProvider)
 	if err != nil {
-		fmt.Println("Token validation failed:", err)
+		return nil, fmt.Errorf("token validation failed: %w", err)
+	}
+
+	if !token.Valid {
+		return nil, fmt.Errorf("invalid token")
+	}
+
+	return rbac.ExtractTokenDetails(token, claims), nil
+}
+
+func checkUserGroups(claims jwt.MapClaims) {
+	groupsInterface, ok := claims["cognito:groups"]
+	if !ok {
+		fmt.Println("No groups found in token.")
 		return
 	}
 
-	// Print token details for debugging
-	rbac.PrintTokenDetails(token)
+	groups, ok := groupsInterface.([]interface{})
+	if !ok {
+		fmt.Println("Invalid groups format in token.")
+		return
+	}
 
-	// Example with Cognito key provider for production
-	/*
-		cognitoProvider := rbac.NewCognitoKeyProvider("us-east-1", "us-east-1_XXXXXXXXX")
-		token, claims, err := rbac.ValidateJWT(tokenString, cognitoProvider)
-	*/
-
-	if token.Valid {
-		// Extract the token details
-		details := rbac.ExtractTokenDetails(token, claims)
-
-		fmt.Println("Token is valid.")
-		fmt.Println("User ID:", details.UserID)
-		fmt.Println("Email:", details.Email)
-		fmt.Println("Expires:", details.Expiration)
-
-		if len(details.Groups) > 0 {
-			fmt.Println("User Groups:")
-			for _, group := range details.Groups {
-				fmt.Println("-", group)
+	if len(groups) > 0 {
+		fmt.Println("User Groups:")
+		for _, group := range groups {
+			if groupStr, ok := group.(string); ok {
+				fmt.Println("-", groupStr)
 			}
-		} else {
-			fmt.Println("No groups found in token.")
-		}
-
-		// Check if user is in a specific group
-		if rbac.HasGroup(claims, "admins") {
-			fmt.Println("User is an admin.")
-		}
-
-		// Check if user is in any of these groups
-		if rbac.HasAnyGroup(claims, []string{"developers", "querybuilders"}) {
-			fmt.Println("User can modify or query data.")
-		}
-
-		// Check if user is in all of these groups
-		if rbac.HasAllGroups(claims, []string{"uploaders", "exporters"}) {
-			fmt.Println("User can both upload and export data.")
 		}
 	} else {
-		fmt.Println("Invalid token")
+		fmt.Println("No groups found in token.")
+	}
+
+	if rbac.HasGroup(claims, "admins") {
+		fmt.Println("User is an admin.")
+	}
+
+	if rbac.HasAnyGroup(claims, []string{"developers", "querybuilders"}) {
+		fmt.Println("User can modify or query data.")
+	}
+
+	if rbac.HasAllGroups(claims, []string{"uploaders", "exporters"}) {
+		fmt.Println("User can both upload and export data.")
 	}
 }
+
+// for testing run like `LOG_LEVEL=INFO go run main.go`
+func main() {
+	_, err := setupConfig()
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	localProvider := rbac.NewLocalKeyProvider("private_key.pem", "public_key.pem")
+
+	details, err := generateAndValidateToken(localProvider)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	fmt.Println("Token is valid.")
+	fmt.Println("User ID:", details.UserID)
+	fmt.Println("Email:", details.Email)
+	fmt.Println("Expires:", details.Expiration)
+
+	checkUserGroups(details.Claims)
+}