fix panic in tests

gravitational · Dec 4, 2024 · e59ddc9 · e59ddc9
1 parent 0240b0f
commit e59ddc9
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 3 deletions.
diff --git 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
@@ -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