From e59ddc908f11a12b87d9b9586327119140481fe9 Mon Sep 17 00:00:00 2001
From: hugoShaka <hugo.hervieux@goteleport.com>
Date: Wed, 4 Dec 2024 17:50:23 -0500
Subject: [PATCH] fix panic in tests

---
 lib/web/autoupdate_common.go      |  4 ++--
 lib/web/autoupdate_common_test.go | 17 ++++++++++++++++-
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/lib/web/autoupdate_common.go b/lib/web/autoupdate_common.go
index 91f7c83554683..1756172f4c6e4 100644
--- a/lib/web/autoupdate_common.go
+++ b/lib/web/autoupdate_common.go
@@ -20,15 +20,15 @@ package web
 
 import (
 	"context"
-	"github.com/gravitational/teleport/api/types"
-	"github.com/gravitational/teleport/lib/utils"
 	"strings"
 
 	"github.com/gravitational/trace"
 
 	autoupdatepb "github.com/gravitational/teleport/api/gen/proto/go/teleport/autoupdate/v1"
+	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/api/types/autoupdate"
 	"github.com/gravitational/teleport/lib/automaticupgrades"
+	"github.com/gravitational/teleport/lib/utils"
 )
 
 // autoUpdateAgentVersion returns the version the agent should install/update to based on
diff --git a/lib/web/autoupdate_common_test.go b/lib/web/autoupdate_common_test.go
index adf2d89365589..a365ac121b078 100644
--- a/lib/web/autoupdate_common_test.go
+++ b/lib/web/autoupdate_common_test.go
@@ -38,6 +38,7 @@ import (
 	"github.com/gravitational/teleport/lib/auth/authclient"
 	"github.com/gravitational/teleport/lib/automaticupgrades"
 	"github.com/gravitational/teleport/lib/automaticupgrades/constants"
+	"github.com/gravitational/teleport/lib/utils"
 )
 
 const (
@@ -167,6 +168,17 @@ func TestAutoUpdateAgentShouldUpdate(t *testing.T) {
 	t.Cleanup(brokenChannelUpstream.Close)
 
 	clock := clockwork.NewFakeClock()
+	cmcCache, err := utils.NewFnCache(utils.FnCacheConfig{
+		TTL:         findEndpointCacheTTL,
+		Clock:       clock,
+		Context:     ctx,
+		ReloadOnErr: false,
+	})
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		cmcCache.Shutdown(ctx)
+	})
+
 	activeUpgradeWindow := types.AgentUpgradeWindow{UTCStartHour: uint32(clock.Now().Hour())}
 	inactiveUpgradeWindow := types.AgentUpgradeWindow{UTCStartHour: uint32(clock.Now().Add(2 * time.Hour).Hour())}
 	tests := []struct {
@@ -309,6 +321,8 @@ func TestAutoUpdateAgentShouldUpdate(t *testing.T) {
 			cmc := types.NewClusterMaintenanceConfig()
 			cmc.SetAgentUpgradeWindow(tt.upgradeWindow)
 			require.NoError(t, tt.channel.CheckAndSetDefaults())
+			// Advance clock to invalidate cache
+			clock.Advance(2 * findEndpointCacheTTL)
 			h := &Handler{
 				cfg: Config{
 					AccessPoint: &fakeRolloutAccessPoint{
@@ -323,7 +337,8 @@ func TestAutoUpdateAgentShouldUpdate(t *testing.T) {
 						groupName: tt.channel,
 					},
 				},
-				clock: clock,
+				clock:                         clock,
+				clusterMaintenanceConfigCache: cmcCache,
 			}
 
 			// Test execution