js/grpc: Tests for multi services in the same proto file

It's an extension of #2265

js/modules/k6/grpc/client.go

@@ -342,7 +342,7 @@ type fileDescriptorLookupKey struct {
 }
 
 func resolveServiceFileDescriptors(
-	mc reflectpb.ServerReflection_ServerReflectionInfoClient,
+	client sendReceiver,
 	res *reflectpb.ListServiceResponse,
 ) (*descriptorpb.FileDescriptorSet, error) {
 	services := res.GetService()
@@ -357,7 +357,7 @@ func resolveServiceFileDescriptors(
 				FileContainingSymbol: service.GetName(),
 			},
 		}
-		resp, err := sendReceive(mc, req)
+		resp, err := sendReceive(client, req)
 		if err != nil {
 			return nil, fmt.Errorf("can't get method on service %q: %w", service, err)
 		}
@@ -384,10 +384,21 @@ func resolveServiceFileDescriptors(
 	return fdset, nil
 }
 
+// sendReceiver is a smaller interface for decoupling
+// from `reflectpb.ServerReflection_ServerReflectionInfoClient`,
+// that has the dependency from `grpc.ClientStream`,
+// which is too much in the case the requirement is to just make a reflection's request.
+// It makes the API more restricted and with a controlled surface,
+// in this way the testing should be easier also.
+type sendReceiver interface {
+	Send(*reflectpb.ServerReflectionRequest) error
+	Recv() (*reflectpb.ServerReflectionResponse, error)
+}
+
 // sendReceive sends a request to a reflection client and,
 // receives a response.
 func sendReceive(
-	client reflectpb.ServerReflection_ServerReflectionInfoClient,
+	client sendReceiver,
 	req *reflectpb.ServerReflectionRequest,
 ) (*reflectpb.ServerReflectionResponse, error) {
 	if err := client.Send(req); err != nil {

js/modules/k6/grpc/client_test.go

@@ -24,13 +24,18 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"fmt"
 	"net/url"
 	"os"
 	"runtime"
 	"strings"
+	"sync/atomic"
 	"testing"
 
 	"google.golang.org/grpc/reflection"
+	reflectpb "google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/descriptorpb"
 
 	"github.com/dop251/goja"
 	"github.com/sirupsen/logrus"
@@ -761,3 +766,103 @@ func TestDebugStat(t *testing.T) {
 		})
 	}
 }
+
+func TestResolveFileDescriptors(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name                string
+		pkgs                []string
+		services            []string
+		expectedDescriptors int
+	}{
+		{
+			name:                "SuccessSamePackage",
+			pkgs:                []string{"mypkg"},
+			services:            []string{"Service1", "Service2", "Service3"},
+			expectedDescriptors: 3,
+		},
+		{
+			name:                "SuccessMultiPackages",
+			pkgs:                []string{"mypkg1", "mypkg2", "mypkg3"},
+			services:            []string{"Service", "Service", "Service"},
+			expectedDescriptors: 3,
+		},
+		{
+			name:                "DeduplicateServices",
+			pkgs:                []string{"mypkg1"},
+			services:            []string{"Service1", "Service2", "Service1"},
+			expectedDescriptors: 2,
+		},
+		{
+			name:                "NoServices",
+			services:            []string{},
+			expectedDescriptors: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			var (
+				lsr  = &reflectpb.ListServiceResponse{}
+				mock = &getServiceFileDescriptorMock{}
+			)
+			for i, service := range tt.services {
+				// if only one package is defined then
+				// the package is the same for every service
+				pkg := tt.pkgs[0]
+				if len(tt.pkgs) > 1 {
+					pkg = tt.pkgs[i]
+				}
+
+				lsr.Service = append(lsr.Service, &reflectpb.ServiceResponse{
+					Name: fmt.Sprintf("%s.%s", pkg, service),
+				})
+				mock.pkgs = append(mock.pkgs, pkg)
+				mock.names = append(mock.names, service)
+			}
+
+			fdset, err := resolveServiceFileDescriptors(mock, lsr)
+			require.NoError(t, err)
+			assert.Len(t, fdset.File, tt.expectedDescriptors)
+		})
+	}
+}
+
+type getServiceFileDescriptorMock struct {
+	nreqs int64
+	pkgs  []string
+	names []string
+}
+
+func (m *getServiceFileDescriptorMock) Send(req *reflectpb.ServerReflectionRequest) error {
+	// TODO: check that the sent message is expected,
+	// otherwise return an error
+	return nil
+}
+
+func (m *getServiceFileDescriptorMock) Recv() (*reflectpb.ServerReflectionResponse, error) {
+	n := atomic.AddInt64(&m.nreqs, 1)
+	ptr := func(s string) (sptr *string) {
+		return &s
+	}
+	index := n - 1
+	fdp := &descriptorpb.FileDescriptorProto{
+		Package: ptr(m.pkgs[index]),
+		Name:    ptr(m.names[index]),
+	}
+	b, err := proto.Marshal(fdp)
+	if err != nil {
+		return nil, err
+	}
+	srr := &reflectpb.ServerReflectionResponse{
+		MessageResponse: &reflectpb.ServerReflectionResponse_FileDescriptorResponse{
+			FileDescriptorResponse: &reflectpb.FileDescriptorResponse{
+				FileDescriptorProto: [][]byte{b},
+			},
+		},
+	}
+	return srr, nil
+}