From 233949e05df3020c4ca58448343352d87f7a4493 Mon Sep 17 00:00:00 2001
From: Jordan Liggitt <liggitt@google.com>
Date: Thu, 30 Nov 2023 14:22:11 -0500
Subject: [PATCH] Add an integration test to verify root path cleanup

---
 .../apiserver/discovery/discovery_test.go     | 13 +++++++---
 .../apiserver/discovery/framework.go          | 25 +++++++++++++++++++
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/test/integration/apiserver/discovery/discovery_test.go b/test/integration/apiserver/discovery/discovery_test.go
index 7ce49612430..71307c1378f 100644
--- a/test/integration/apiserver/discovery/discovery_test.go
+++ b/test/integration/apiserver/discovery/discovery_test.go
@@ -37,6 +37,7 @@ import (
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	runtimeserializer "k8s.io/apimachinery/pkg/runtime/serializer"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
+	"k8s.io/apimachinery/pkg/util/sets"
 	discoveryendpoint "k8s.io/apiserver/pkg/endpoints/discovery/aggregated"
 	genericfeatures "k8s.io/apiserver/pkg/features"
 	utilfeature "k8s.io/apiserver/pkg/util/feature"
@@ -241,6 +242,7 @@ func TestAggregatedAPIServiceDiscovery(t *testing.T) {
 
 	// For each groupversion served by our resourcemanager, create an APIService
 	// object connected to our fake APIServer
+	var groupVersions []metav1.GroupVersion
 	for _, versionInfo := range basicTestGroup.Versions {
 		groupVersion := metav1.GroupVersion{
 			Group:   basicTestGroup.Name,
@@ -248,14 +250,19 @@ func TestAggregatedAPIServiceDiscovery(t *testing.T) {
 		}
 
 		require.NoError(t, registerAPIService(ctx, client, groupVersion, service))
-		defer func() {
-			require.NoError(t, unregisterAPIService(ctx, client, groupVersion))
-		}()
+		groupVersions = append(groupVersions, groupVersion)
 	}
 
 	// Keep repeatedly fetching document from aggregator.
 	// Check to see if it contains our service within a reasonable amount of time
 	require.NoError(t, WaitForGroups(ctx, client, basicTestGroupWithFixup))
+	require.NoError(t, WaitForRootPaths(t, ctx, client, sets.New("/apis/"+basicTestGroup.Name), nil))
+
+	// Unregister and ensure the group gets dropped from root paths
+	for _, groupVersion := range groupVersions {
+		require.NoError(t, unregisterAPIService(ctx, client, groupVersion))
+	}
+	require.NoError(t, WaitForRootPaths(t, ctx, client, nil, sets.New("/apis/"+basicTestGroup.Name)))
 }
 
 func runTestCases(t *testing.T, cases []testCase) {
diff --git a/test/integration/apiserver/discovery/framework.go b/test/integration/apiserver/discovery/framework.go
index 4593b27bb9e..f0642330db9 100644
--- a/test/integration/apiserver/discovery/framework.go
+++ b/test/integration/apiserver/discovery/framework.go
@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"reflect"
 	"strings"
+	"testing"
 	"time"
 
 	apiextensions "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
@@ -29,6 +30,7 @@ import (
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/apimachinery/pkg/util/wait"
 	"k8s.io/client-go/dynamic"
 	"k8s.io/client-go/kubernetes"
@@ -573,6 +575,29 @@ func WaitForGroupsAbsent(ctx context.Context, client testClient, groups ...strin
 
 }
 
+func WaitForRootPaths(t *testing.T, ctx context.Context, client testClient, requirePaths, forbidPaths sets.Set[string]) error {
+	return wait.PollUntilContextTimeout(ctx, 250*time.Millisecond, maxTimeout, true, func(ctx context.Context) (done bool, err error) {
+		statusContent, err := client.Discovery().RESTClient().Get().AbsPath("/").SetHeader("Accept", "application/json").DoRaw(ctx)
+		if err != nil {
+			return false, err
+		}
+		rootPaths := metav1.RootPaths{}
+		if err := json.Unmarshal(statusContent, &rootPaths); err != nil {
+			return false, err
+		}
+		paths := sets.New(rootPaths.Paths...)
+		if missing := requirePaths.Difference(paths); len(missing) > 0 {
+			t.Logf("missing required root paths %v", sets.List(missing))
+			return false, nil
+		}
+		if present := forbidPaths.Intersection(paths); len(present) > 0 {
+			t.Logf("present forbidden root paths %v", sets.List(present))
+			return false, nil
+		}
+		return true, nil
+	})
+}
+
 func WaitForGroups(ctx context.Context, client testClient, groups ...apidiscoveryv2beta1.APIGroupDiscovery) error {
 	return WaitForResultWithCondition(ctx, client, func(groupList apidiscoveryv2beta1.APIGroupDiscoveryList) bool {
 		for _, searchGroup := range groups {