Revert triton_test.go and kuma_test.go to main version

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

discovery/triton/triton_test.go

@@ -80,7 +80,7 @@ var (
 	}
 )
 
-func newTritonDiscovery(c SDConfig) (*Discovery, func(), error) {
+func newTritonDiscovery(c SDConfig) (*Discovery, discovery.DiscovererMetrics, error) {
 	reg := prometheus.NewRegistry()
 	refreshMetrics := discovery.NewRefreshMetrics(reg)
 	metrics := c.NewDiscovererMetrics(reg, refreshMetrics)
@@ -98,16 +98,11 @@ func newTritonDiscovery(c SDConfig) (*Discovery, func(), error) {
 		return nil, nil, err
 	}
 
-	cleanup := func() {
+	return d, metrics, nil
-		metrics.Unregister()
-		refreshMetrics.Unregister()
-	}
-
-	return d, cleanup, nil
 }
 
 func TestTritonSDNew(t *testing.T) {
-	td, cleanup, err := newTritonDiscovery(conf)
+	td, m, err := newTritonDiscovery(conf)
 	require.NoError(t, err)
 	require.NotNil(t, td)
 	require.NotNil(t, td.client)
@@ -117,20 +112,17 @@ func TestTritonSDNew(t *testing.T) {
 	require.Equal(t, conf.DNSSuffix, td.sdConfig.DNSSuffix)
 	require.Equal(t, conf.Endpoint, td.sdConfig.Endpoint)
 	require.Equal(t, conf.Port, td.sdConfig.Port)
-	defer cleanup()
+	m.Unregister()
 }
 
 func TestTritonSDNewBadConfig(t *testing.T) {
-	td, cleanup, err := newTritonDiscovery(badconf)
+	td, _, err := newTritonDiscovery(badconf)
 	require.Error(t, err)
 	require.Nil(t, td)
-	if cleanup != nil {
-		defer cleanup()
-	}
 }
 
 func TestTritonSDNewGroupsConfig(t *testing.T) {
-	td, cleanup, err := newTritonDiscovery(groupsconf)
+	td, m, err := newTritonDiscovery(groupsconf)
 	require.NoError(t, err)
 	require.NotNil(t, td)
 	require.NotNil(t, td.client)
@@ -141,11 +133,11 @@ func TestTritonSDNewGroupsConfig(t *testing.T) {
 	require.Equal(t, groupsconf.Endpoint, td.sdConfig.Endpoint)
 	require.Equal(t, groupsconf.Groups, td.sdConfig.Groups)
 	require.Equal(t, groupsconf.Port, td.sdConfig.Port)
-	defer cleanup()
+	m.Unregister()
 }
 
 func TestTritonSDNewCNConfig(t *testing.T) {
-	td, cleanup, err := newTritonDiscovery(cnconf)
+	td, m, err := newTritonDiscovery(cnconf)
 	require.NoError(t, err)
 	require.NotNil(t, td)
 	require.NotNil(t, td.client)
@@ -156,7 +148,7 @@ func TestTritonSDNewCNConfig(t *testing.T) {
 	require.Equal(t, cnconf.DNSSuffix, td.sdConfig.DNSSuffix)
 	require.Equal(t, cnconf.Endpoint, td.sdConfig.Endpoint)
 	require.Equal(t, cnconf.Port, td.sdConfig.Port)
-	defer cleanup()
+	m.Unregister()
 }
 
 func TestTritonSDRefreshNoTargets(t *testing.T) {
@@ -189,21 +181,21 @@ func TestTritonSDRefreshMultipleTargets(t *testing.T) {
 }
 
 func TestTritonSDRefreshNoServer(t *testing.T) {
-	td, cleanup, _ := newTritonDiscovery(conf)
+	td, m, _ := newTritonDiscovery(conf)
-	defer cleanup()
 
 	_, err := td.refresh(context.Background())
 	require.ErrorContains(t, err, "an error occurred when requesting targets from the discovery endpoint")
+	m.Unregister()
 }
 
 func TestTritonSDRefreshCancelled(t *testing.T) {
-	td, cleanup, _ := newTritonDiscovery(conf)
+	td, m, _ := newTritonDiscovery(conf)
-	defer cleanup()
 
 	ctx, cancel := context.WithCancel(context.Background())
 	cancel()
 	_, err := td.refresh(ctx)
 	require.ErrorContains(t, err, context.Canceled.Error())
+	m.Unregister()
 }
 
 func TestTritonSDRefreshCNsUUIDOnly(t *testing.T) {
@@ -239,12 +231,12 @@ func TestTritonSDRefreshCNsWithHostname(t *testing.T) {
 }
 
 func testTritonSDRefresh(t *testing.T, c SDConfig, dstr string) []model.LabelSet {
-	td, cleanup, _ := newTritonDiscovery(c)
+	var (
-	defer cleanup()
+		td, m, _ = newTritonDiscovery(c)
-
+		s        = httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
-	s := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			fmt.Fprintln(w, dstr)
-		fmt.Fprintln(w, dstr)
+		}))
-	}))
+	)
 
 	defer s.Close()
 
@@ -269,5 +261,7 @@ func testTritonSDRefresh(t *testing.T, c SDConfig, dstr string) []model.LabelSet
 	tg := tgs[0]
 	require.NotNil(t, tg)
 
+	m.Unregister()
+
 	return tg.Targets
 }

discovery/xds/kuma_test.go

@@ -108,31 +108,25 @@ func getKumaMadsV1DiscoveryResponse(resources ...*MonitoringAssignment) (*v3.Dis
 	}, nil
 }
 
-func newKumaTestHTTPDiscovery(c KumaSDConfig) (*fetchDiscovery, func(), error) {
+func newKumaTestHTTPDiscovery(c KumaSDConfig) (*fetchDiscovery, error) {
 	reg := prometheus.NewRegistry()
 	refreshMetrics := discovery.NewRefreshMetrics(reg)
 	metrics := c.NewDiscovererMetrics(reg, refreshMetrics)
 	err := metrics.Register()
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
 	kd, err := NewKumaHTTPDiscovery(&c, nopLogger, metrics)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
 	pd, ok := kd.(*fetchDiscovery)
 	if !ok {
-		return nil, nil, errors.New("not a fetchDiscovery")
+		return nil, errors.New("not a fetchDiscovery")
 	}
-
+	return pd, nil
-	cleanup := func() {
-		metrics.Unregister()
-		refreshMetrics.Unregister()
-	}
-
-	return pd, cleanup, nil
 }
 
 func TestKumaMadsV1ResourceParserInvalidTypeURL(t *testing.T) {
@@ -221,10 +215,9 @@ func TestKumaMadsV1ResourceParserInvalidResources(t *testing.T) {
 func TestNewKumaHTTPDiscovery(t *testing.T) {
 	t.Parallel()
 
-	kd, cleanup, err := newKumaTestHTTPDiscovery(kumaConf)
+	kd, err := newKumaTestHTTPDiscovery(kumaConf)
 	require.NoError(t, err)
 	require.NotNil(t, kd)
-	defer cleanup()
 
 	resClient, ok := kd.client.(*HTTPResourceClient)
 	require.True(t, ok)
@@ -232,6 +225,8 @@ 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()
 }
 
 func TestKumaHTTPDiscoveryRefresh(t *testing.T) {
@@ -263,10 +258,9 @@ tls_config:
 	var cfg KumaSDConfig
 	require.NoError(t, yaml.Unmarshal([]byte(cfgString), &cfg))
 
-	kd, cleanup, err := newKumaTestHTTPDiscovery(cfg)
+	kd, err := newKumaTestHTTPDiscovery(cfg)
 	require.NoError(t, err)
 	require.NotNil(t, kd)
-	defer cleanup()
 
 	ch := make(chan []*targetgroup.Group, 1)
 	kd.poll(context.Background(), ch)
@@ -329,4 +323,6 @@ tls_config:
 	case <-ch:
 		require.Fail(t, "no update expected")
 	}
+
+	kd.metrics.Unregister()
 }