chore: Restrict network-dependent helper to integration tests (#4250)

Author: callpraths

pkg/experiment/metastore/client/server_mock_test.go

@@ -3,7 +3,6 @@ package metastoreclient
 import (
 	"context"
 	"fmt"
-	"net"
 	"sync"
 	"testing"
 
@@ -19,6 +18,7 @@ import (
 	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
 	"github.com/grafana/pyroscope/pkg/experiment/metastore/discovery"
 	"github.com/grafana/pyroscope/pkg/experiment/metastore/raftnode/raftnodepb"
+	"github.com/grafana/pyroscope/pkg/test"
 	"github.com/grafana/pyroscope/pkg/test/mocks/mockmetastorev1"
 	"github.com/grafana/pyroscope/pkg/test/mocks/mockraftnodepb"
 )
@@ -201,16 +201,14 @@ func errOrT[T any](t *T, f func() error) (*T, error) {
 // Returns the grpc.DialOptions needed for a client connection to the created mock servers.
 func createMockServers(t *testing.T, l log.Logger, dServers []discovery.Server) (*mockServers, []grpc.DialOption) {
 	var servers []*mockServer
-	listeners := make(map[string]*bufconn.Listener)
+	listeners, dialOpt := createInMemoryListeners(dServers)
 	for idx, dserv := range dServers {
 		s := newMockServer(t)
 		s.index = idx
 		s.id = testServerId(idx)
 		s.address = dserv.ResolvedAddress
-		lis := bufconn.Listen(256 << 10)
-		listeners[s.address] = lis
 		go func() {
-			if err := s.srv.Serve(lis); err != nil {
+			if err := s.srv.Serve(listeners[s.address]); err != nil {
 				assert.NoError(t, err)
 			}
 		}()
@@ -222,14 +220,15 @@ func createMockServers(t *testing.T, l log.Logger, dServers []discovery.Server)
 		t:       t,
 		l:       l,
 	}
-	dialer := func(_ context.Context, address string) (net.Conn, error) {
-		el := listeners[address]
-		if el != nil {
-			return el.Dial()
-		}
-		return net.Dial("tcp", address)
+	return ms, []grpc.DialOption{dialOpt}
+}
+
+func createInMemoryListeners(dServers []discovery.Server) (map[string]*bufconn.Listener, grpc.DialOption) {
+	addrs := make([]string, 0, len(dServers))
+	for _, ds := range dServers {
+		addrs = append(addrs, ds.ResolvedAddress)
 	}
-	return ms, []grpc.DialOption{grpc.WithContextDialer(dialer)}
+	return test.CreateInMemoryListeners(addrs)
 }
 
 func newMockServer(t *testing.T) *mockServer {

pkg/experiment/metastore/test/create.go

@@ -3,7 +3,6 @@ package test
 import (
 	"context"
 	"fmt"
-	"net"
 	"testing"
 	"time"
 
@@ -32,24 +31,20 @@ import (
 func NewMetastoreSet(t *testing.T, cfg *metastore.Config, n int, bucket objstore.Bucket) MetastoreSet {
 	l := test.NewTestingLogger(t)
 
-	ports, err := test.GetFreePorts(2 * n)
-	addresses := make([]string, 2*n)
-	for i, port := range ports {
-		addresses[i] = fmt.Sprintf("localhost:%d", port)
-	}
-	grpcAddresses := addresses[:n]
-	raftAddresses := addresses[n:]
+	grpcAddresses := make([]string, n)
+	raftAddresses := make([]string, n)
 	raftIds := make([]string, n)
+	bootstrapPeers := make([]string, n)
 	for i := 0; i < n; i++ {
+		grpcAddresses[i] = fmt.Sprintf("localhost:%d", 10500+i)
+		raftAddresses[i] = fmt.Sprintf("localhost:%d", 10500+2*i)
 		raftIds[i] = fmt.Sprintf("node-%d", i)
+		bootstrapPeers[i] = fmt.Sprintf("%s/%s", raftAddresses[i], raftIds[i])
 	}
-	bootstrapPeers := make([]string, n)
-	configs := make([]metastore.Config, n)
-	servers := make([]discovery.Server, n)
+	l.Log("grpcAddresses", fmt.Sprintf("%+v", grpcAddresses), "raftAddresses", fmt.Sprintf("%+v", raftAddresses))
 
+	configs := make([]metastore.Config, n)
 	for i := 0; i < n; i++ {
-		bootstrapPeers[i] = fmt.Sprintf("%s/%s", raftAddresses[i], raftIds[i])
-
 		icfg := *cfg
 		icfg.MinReadyDuration = 0
 		icfg.Address = grpcAddresses[i]
@@ -60,30 +55,35 @@ func NewMetastoreSet(t *testing.T, cfg *metastore.Config, n int, bucket objstore
 		icfg.Raft.BindAddress = raftAddresses[i]
 		icfg.Raft.BootstrapPeers = bootstrapPeers
 		icfg.Raft.BootstrapExpectPeers = n
+		configs[i] = icfg
+	}
+
+	servers := make([]discovery.Server, n)
+	for i := 0; i < n; i++ {
 		srv := discovery.Server{
 			Raft: raft.Server{
 				ID:      raft.ServerID(raftIds[i]),
 				Address: raft.ServerAddress(raftAddresses[i]),
 			},
-			ResolvedAddress: addresses[i],
+			ResolvedAddress: grpcAddresses[i],
 		}
 		servers[i] = srv
-		configs[i] = icfg
 	}
-	require.NoError(t, err)
 
+	listeners, dialOpt := test.CreateInMemoryListeners(grpcAddresses)
 	d := MockStaticDiscovery(t, servers)
-	client := metastoreclient.New(l, cfg.GRPCClientConfig, d)
-	err = client.Service().StartAsync(context.Background())
+	client := metastoreclient.New(l, cfg.GRPCClientConfig, d, dialOpt)
+	err := client.Service().StartAsync(context.Background())
 	require.NoError(t, err)
 
-	l.Log("grpcAddresses", fmt.Sprintf("%+v", grpcAddresses), "raftAddresses", fmt.Sprintf("%+v", raftAddresses))
 	res := MetastoreSet{
 		t: t,
 	}
+
 	for i := 0; i < n; i++ {
 		options, err := cfg.GRPCClientConfig.DialOption(nil, nil)
 		require.NoError(t, err)
+		options = append(options, dialOpt)
 		cc, err := grpc.Dial(grpcAddresses[i], options...)
 		require.NoError(t, err)
 		logger := log.With(l, "idx", bootstrapPeers[i])
@@ -100,8 +100,7 @@ func NewMetastoreSet(t *testing.T, cfg *metastore.Config, n int, bucket objstore
 		require.NoError(t, err)
 		m.Register(server)
 
-		lis, err := net.Listen("tcp", grpcAddresses[i])
-		assert.NoError(t, err)
+		lis := listeners[grpcAddresses[i]]
 		go func() {
 			err := server.Serve(lis)
 			assert.NoError(t, err)

pkg/experiment/query_backend/client/client_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"flag"
 	"fmt"
-	"net"
 	"testing"
 	"time"
 
@@ -14,7 +13,6 @@ import (
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/resolver"
-	"google.golang.org/grpc/test/bufconn"
 
 	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
 	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
@@ -83,28 +81,20 @@ func (r *multiResolver) Close() {}
 // actual network I/O.
 func Test_Concurrency(t *testing.T) {
 	addresses := make([]string, 0, nServers)
-	listeners := make(map[string]*bufconn.Listener)
-
-	// Create a separate bufconn listener to be used by each server setup below.
 	for i := 0; i < nServers; i++ {
 		address := fmt.Sprintf("localhost:%d", 10004+i)
 		addresses = append(addresses, address)
-		listeners[address] = bufconn.Listen(256 << 10)
-	}
-	// In-memory dialer for the client dials the appropriate bufconn listener.
-	dialer := func(_ context.Context, address string) (net.Conn, error) {
-		listener := listeners[address]
-		require.NotNil(t, listener)
-		return listener.Dial()
 	}
 
+	listeners, dialOpt := test.CreateInMemoryListeners(addresses)
+
 	grpcClientCfg := grpcclient.Config{}
 	grpcClientCfg.RegisterFlags(flag.NewFlagSet("", flag.PanicOnError))
 
 	resolver.Register(&multiResolverBuilder{targets: addresses})
 	backendAddress := "multi:///"
 
-	cl, err := New(backendAddress, grpcClientCfg, grpc.WithContextDialer(dialer))
+	cl, err := New(backendAddress, grpcClientCfg, dialOpt)
 	require.NoError(t, err)
 
 	for i := 0; i < nServers; i++ {

pkg/test/in_memory_listeners.go

@@ -0,0 +1,28 @@
+package test
+
+import (
+	"context"
+	"net"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/test/bufconn"
+)
+
+// Create in-memory listeners at given addresses.
+// Also returns gRPC dial option for a client to connect to the appropriate in-memory listener
+// for a given address.
+func CreateInMemoryListeners(addresses []string) (map[string]*bufconn.Listener, grpc.DialOption) {
+	listeners := make(map[string]*bufconn.Listener)
+	for _, a := range addresses {
+		el := bufconn.Listen(256 << 10)
+		listeners[a] = el
+	}
+	dialer := func(_ context.Context, address string) (net.Conn, error) {
+		el := listeners[address]
+		if el != nil {
+			return el.Dial()
+		}
+		return net.Dial("tcp", address)
+	}
+	return listeners, grpc.WithContextDialer(dialer)
+}

pkg/test/integration/helper.go

@@ -17,8 +17,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/grafana/pyroscope/pkg/test"
-
 	"connectrpc.com/connect"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/assert"
@@ -109,7 +107,7 @@ func (p *PyroscopeTest) start(t *testing.T) {
 }
 
 func (p *PyroscopeTest) Configure(t *testing.T, v2 bool) *PyroscopeTest {
-	ports, err := test.GetFreePorts(4)
+	ports, err := GetFreePorts(4)
 	require.NoError(t, err)
 	p.httpPort = ports[0]
 	p.memberlistPort = ports[1]

pkg/test/ports.go renamed to pkg/test/integration/ports.go

@@ -1,8 +1,12 @@
-package test
+package integration
 
 import "net"
 
-// GetFreePorts returns a number of free local port for the tests to listen on. Note this will make sure the returned ports do not overlap, by stopping to listen once all ports are allocated
+// GetFreePorts returns a number of free local port for the tests to listen on.
+// This will make sure the returned ports do not overlap, by stopping to listen once all ports are allocated
+//
+// Note: This function should only be used in integration tests.
+// Use in-memory network connections in unittests.
 func GetFreePorts(len int) (ports []int, err error) {
 	ports = make([]int, len)
 	for i := 0; i < len; i++ {