From e04a6fb08e3b994ee63cd80503630b1937807d91 Mon Sep 17 00:00:00 2001
From: Jo <joao.guerreiro@grafana.com>
Date: Mon, 31 Jul 2023 10:41:37 +0200
Subject: [PATCH] Auth: Fix static test fake (#72514)

fix static fake
---
 pkg/services/authn/authntest/fake.go          | 53 +++++++++++++++----
 pkg/services/contexthandler/contexthandler.go |  6 +--
 2 files changed, 47 insertions(+), 12 deletions(-)

diff --git a/pkg/services/authn/authntest/fake.go b/pkg/services/authn/authntest/fake.go
index 878c9b926fa..b28e349c05f 100644
--- a/pkg/services/authn/authntest/fake.go
+++ b/pkg/services/authn/authntest/fake.go
@@ -9,28 +9,63 @@ import (
 var _ authn.Service = new(FakeService)
 
 type FakeService struct {
-	ExpectedErr      error
-	ExpectedRedirect *authn.Redirect
-	ExpectedIdentity *authn.Identity
+	ExpectedErr        error
+	ExpectedRedirect   *authn.Redirect
+	ExpectedIdentity   *authn.Identity
+	ExpectedErrs       []error
+	ExpectedIdentities []*authn.Identity
+	CurrentIndex       int
 }
 
-func (f FakeService) Authenticate(ctx context.Context, r *authn.Request) (*authn.Identity, error) {
+func (f *FakeService) Authenticate(ctx context.Context, r *authn.Request) (*authn.Identity, error) {
+	if f.ExpectedIdentities != nil {
+		if f.CurrentIndex >= len(f.ExpectedIdentities) {
+			panic("ExpectedIdentities is empty")
+		}
+		if f.CurrentIndex >= len(f.ExpectedErrs) {
+			panic("ExpectedErrs is empty")
+		}
+
+		identity := f.ExpectedIdentities[f.CurrentIndex]
+		err := f.ExpectedErrs[f.CurrentIndex]
+
+		f.CurrentIndex += 1
+
+		return identity, err
+	}
+
 	return f.ExpectedIdentity, f.ExpectedErr
 }
 
-func (f FakeService) RegisterPostAuthHook(hook authn.PostAuthHookFn, priority uint) {}
+func (f *FakeService) RegisterPostAuthHook(hook authn.PostAuthHookFn, priority uint) {}
+
+func (f *FakeService) Login(ctx context.Context, client string, r *authn.Request) (*authn.Identity, error) {
+	if f.ExpectedIdentities != nil {
+		if f.CurrentIndex >= len(f.ExpectedIdentities) {
+			panic("ExpectedIdentities is empty")
+		}
+		if f.CurrentIndex >= len(f.ExpectedErrs) {
+			panic("ExpectedErrs is empty")
+		}
+
+		identity := f.ExpectedIdentities[f.CurrentIndex]
+		err := f.ExpectedErrs[f.CurrentIndex]
+
+		f.CurrentIndex += 1
+
+		return identity, err
+	}
 
-func (f FakeService) Login(ctx context.Context, client string, r *authn.Request) (*authn.Identity, error) {
 	return f.ExpectedIdentity, f.ExpectedErr
 }
 
-func (f FakeService) RegisterPostLoginHook(hook authn.PostLoginHookFn, priority uint) {}
+func (f *FakeService) RegisterPostLoginHook(hook authn.PostLoginHookFn, priority uint) {}
 
-func (f FakeService) RedirectURL(ctx context.Context, client string, r *authn.Request) (*authn.Redirect, error) {
+func (f *FakeService) RedirectURL(ctx context.Context, client string, r *authn.Request) (*authn.Redirect, error) {
 	return f.ExpectedRedirect, f.ExpectedErr
 }
 
-func (f FakeService) RegisterClient(c authn.Client) {}
+func (f *FakeService) RegisterClient(c authn.Client) {}
 
 var _ authn.ContextAwareClient = new(FakeClient)
 
diff --git a/pkg/services/contexthandler/contexthandler.go b/pkg/services/contexthandler/contexthandler.go
index a0b664cf8bf..fdeaf39da2f 100644
--- a/pkg/services/contexthandler/contexthandler.go
+++ b/pkg/services/contexthandler/contexthandler.go
@@ -70,7 +70,7 @@ func ProvideService(cfg *setting.Cfg, tokenService auth.UserTokenService, jwtSer
 		orgService:        orgService,
 		oauthTokenService: oauthTokenService,
 		features:          features,
-		authnService:      authnService,
+		AuthnService:      authnService,
 		anonDeviceService: anonDeviceService,
 		singleflight:      new(singleflight.Group),
 	}
@@ -93,7 +93,7 @@ type ContextHandler struct {
 	orgService        org.Service
 	oauthTokenService oauthtoken.OAuthTokenService
 	features          *featuremgmt.FeatureManager
-	authnService      authn.Service
+	AuthnService      authn.Service
 	singleflight      *singleflight.Group
 	anonDeviceService anonymous.Service
 	// GetTime returns the current time.
@@ -171,7 +171,7 @@ func (h *ContextHandler) Middleware(next http.Handler) http.Handler {
 		}
 
 		if h.Cfg.AuthBrokerEnabled {
-			identity, err := h.authnService.Authenticate(ctx, &authn.Request{HTTPRequest: reqContext.Req, Resp: reqContext.Resp})
+			identity, err := h.AuthnService.Authenticate(ctx, &authn.Request{HTTPRequest: reqContext.Req, Resp: reqContext.Resp})
 			if err != nil {
 				if errors.Is(err, auth.ErrInvalidSessionToken) {
 					// Burn the cookie in case of invalid, expired or missing token