Added cleanup closure in kuma test

Signed-off-by: intojhanurag <aojharaj2004@gmail.com>

discovery/xds/kuma_test.go

@@ -108,7 +108,7 @@ func getKumaMadsV1DiscoveryResponse(resources ...*MonitoringAssignment) (*v3.Dis
 	}, nil
 }
 
-func newKumaTestHTTPDiscovery(c KumaSDConfig) (*fetchDiscovery, discovery.RefreshMetricsManager, error) {
+func newKumaTestHTTPDiscovery(c KumaSDConfig) (*fetchDiscovery, func(), error) {
 	reg := prometheus.NewRegistry()
 	refreshMetrics := discovery.NewRefreshMetrics(reg)
 	metrics := c.NewDiscovererMetrics(reg, refreshMetrics)
@@ -126,7 +126,13 @@ func newKumaTestHTTPDiscovery(c KumaSDConfig) (*fetchDiscovery, discovery.Refres
 	if !ok {
 		return nil, nil, errors.New("not a fetchDiscovery")
 	}
-	return pd, refreshMetrics, nil
+
+	cleanup := func() {
+		metrics.Unregister()
+		refreshMetrics.Unregister()
+	}
+
+	return pd, cleanup, nil
 }
 
 func TestKumaMadsV1ResourceParserInvalidTypeURL(t *testing.T) {
@@ -215,9 +221,10 @@ func TestKumaMadsV1ResourceParserInvalidResources(t *testing.T) {
 func TestNewKumaHTTPDiscovery(t *testing.T) {
 	t.Parallel()
 
-	kd, rm, err := newKumaTestHTTPDiscovery(kumaConf)
+	kd, cleanup, err := newKumaTestHTTPDiscovery(kumaConf)
 	require.NoError(t, err)
 	require.NotNil(t, kd)
+	defer cleanup()
 
 	resClient, ok := kd.client.(*HTTPResourceClient)
 	require.True(t, ok)
@@ -225,9 +232,6 @@ func TestNewKumaHTTPDiscovery(t *testing.T) {
 	require.Equal(t, KumaMadsV1ResourceTypeURL, resClient.ResourceTypeURL())
 	require.Equal(t, kumaConf.ClientID, resClient.ID())
 	require.Equal(t, KumaMadsV1ResourceType, resClient.config.ResourceType)
-
-	kd.metrics.Unregister()
-	rm.Unregister()
 }
 
 func TestKumaHTTPDiscoveryRefresh(t *testing.T) {
@@ -259,9 +263,10 @@ tls_config:
 	var cfg KumaSDConfig
 	require.NoError(t, yaml.Unmarshal([]byte(cfgString), &cfg))
 
-	kd, rm, err := newKumaTestHTTPDiscovery(cfg)
+	kd, cleanup, err := newKumaTestHTTPDiscovery(cfg)
 	require.NoError(t, err)
 	require.NotNil(t, kd)
+	defer cleanup()
 
 	ch := make(chan []*targetgroup.Group, 1)
 	kd.poll(context.Background(), ch)
@@ -320,13 +325,8 @@ tls_config:
 	kd.poll(ctx, ch)
 	select {
 	case <-ctx.Done():
-		kd.metrics.Unregister()
-		rm.Unregister()
 		return
 	case <-ch:
 		require.Fail(t, "no update expected")
 	}
-
-	kd.metrics.Unregister()
-	rm.Unregister()
 }