From 78024dea74470d48805cc24e6eb043c5a86bf29e Mon Sep 17 00:00:00 2001
From: Yoriyasu Yano <430092+yorinasub17@users.noreply.github.com>
Date: Sat, 11 Mar 2023 06:59:50 -0600
Subject: [PATCH] Allow overriding the default user that is returned when the
 queue is empty

Signed-off-by: Yoriyasu Yano <430092+yorinasub17@users.noreply.github.com>
---
 README.md   | 11 ++++++++++-
 mockoidc.go | 10 +++++++++-
 queue.go    | 14 ++++++++++++--
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 81669d4..672cef7 100644
--- a/README.md
+++ b/README.md
@@ -32,7 +32,7 @@ cfg := m.Config()
 //   	ClientID     string
 //   	ClientSecret string
 //   	Issuer       string
-//   
+//
 //   	AccessTTL  time.Duration
 //   	RefreshTTL time.Duration
 // }
@@ -97,6 +97,15 @@ m.QueueCode("12345")
 // ...Request to m.AuthorizationEndpoint()
 ```
 
+You can also override the default user that is returned:
+
+```
+user := &mockoidc.User{
+    // User details...
+}
+m.SetDefaultUser(user)
+```
+
 ### Forcing Errors
 
 Arbitrary errors can also be queued for handlers to return instead of their
diff --git a/mockoidc.go b/mockoidc.go
index 990923f..d55c28d 100644
--- a/mockoidc.go
+++ b/mockoidc.go
@@ -79,7 +79,7 @@ func NewServer(key *rsa.PrivateKey) (*MockOIDC, error) {
 		CodeChallengeMethodsSupported: []string{"plain", "S256"},
 		Keypair:                       keypair,
 		SessionStore:                  NewSessionStore(),
-		UserQueue:                     &UserQueue{},
+		UserQueue:                     &UserQueue{DefaultUser: DefaultUser()},
 		ErrorQueue:                    &ErrorQueue{},
 	}, nil
 }
@@ -169,6 +169,14 @@ func (m *MockOIDC) QueueUser(user User) {
 	m.UserQueue.Push(user)
 }
 
+// SetDefaultUser allows setting the default mock User object to return for
+// authentication when the authentication queue is empty.
+// Calls to the `authorization_endpoint` will return the set user instead of
+// `DefaultUser()` when the queue is empty.
+func (m *MockOIDC) SetDefaultUser(user User) {
+	m.UserQueue.SetDefaultUser(user)
+}
+
 // QueueCode allows adding mock code strings to the authentication queue.
 // Calls to the `authorization_endpoint` will pop these code strings
 // off the queue and create a session with them and return them as the
diff --git a/queue.go b/queue.go
index bd706d3..d989ae6 100644
--- a/queue.go
+++ b/queue.go
@@ -7,6 +7,8 @@ import "sync"
 type UserQueue struct {
 	sync.Mutex
 	Queue []User
+
+	DefaultUser User
 }
 
 // CodeQueue manages the queue of codes returned for each
@@ -37,12 +39,15 @@ func (q *UserQueue) Push(user User) {
 	q.Queue = append(q.Queue, user)
 }
 
-// Pop a User from the Queue. If empty, return `DefaultUser()`
+// Pop a User from the Queue. If empty, return `q.DefaultUser` if set, otherwise
+// `DefaultUser()`
 func (q *UserQueue) Pop() User {
 	q.Lock()
 	defer q.Unlock()
 
-	if len(q.Queue) == 0 {
+	if len(q.Queue) == 0 && q.DefaultUser != nil {
+		return q.DefaultUser
+	} else if len(q.Queue) == 0 {
 		return DefaultUser()
 	}
 
@@ -51,6 +56,11 @@ func (q *UserQueue) Pop() User {
 	return user
 }
 
+// SetDefaultUser configures the queue with the default user to return when empty.
+func (q *UserQueue) SetDefaultUser(user User) {
+	q.DefaultUser = user
+}
+
 // Push adds a code to the Queue to be returned by subsequent
 // `authorization_endpoint` calls as the code
 func (q *CodeQueue) Push(code string) {