fix: make mock package more flexible and support wrapperspb

This replaces use of the mock.Register() test helper function with
a more flexible solution that supports wrapperspb types. The only
caveat is that this requires pre-registration via TestMain, but
is much more flexible, since we can add custom services, methods
and message types for our testing needs without having to define
them via proto files that needs to be generated.

Note: this adds mock.GetValue, which will be used in a future test.

Author: meling

.vscode/gorums.txt

@@ -97,5 +97,6 @@ unmarshaled
 unmarshaling
 unmarshals
 userguide
+wrapperspb
 Xeon
 Zorums

channel_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/relab/gorums/ordering"
 	"google.golang.org/grpc"
 	"google.golang.org/protobuf/proto"
+	pb "google.golang.org/protobuf/types/known/wrapperspb"
 )
 
 const defaultTestTimeout = 3 * time.Second
@@ -34,15 +35,15 @@ func newNodeWithServer(t testing.TB, delay time.Duration) *RawNode {
 type mockSrv struct{}
 
 func (mockSrv) Test(_ ServerCtx, req proto.Message) (proto.Message, error) {
-	return mock.NewResponse(mock.GetVal(req) + "-mocked-"), nil
+	return pb.String(mock.GetVal(req) + "-mocked-"), nil
 }
 
 func newNodeWithStoppableServer(t testing.TB, delay time.Duration) (*RawNode, func()) {
 	t.Helper()
 	addrs, teardown := TestSetup(t, 1, func(_ int) ServerIface {
 		mockSrv := &mockSrv{}
 		srv := NewServer()
-		srv.RegisterHandler(mock.ServerMethodName, func(ctx ServerCtx, in *Message) (*Message, error) {
+		srv.RegisterHandler(mock.TestMethod, func(ctx ServerCtx, in *Message) (*Message, error) {
 			// Simulate slow processing
 			time.Sleep(delay)
 			resp, err := mockSrv.Test(ctx, in.GetProtoMessage())
@@ -59,7 +60,7 @@ func sendRequest(t testing.TB, node *RawNode, req request, msgID uint64) respons
 	if req.ctx == nil {
 		req.ctx = t.Context()
 	}
-	req.msg = NewRequestMessage(ordering.NewGorumsMetadata(req.ctx, msgID, mock.ServerMethodName), nil)
+	req.msg = NewRequestMessage(ordering.NewGorumsMetadata(req.ctx, msgID, mock.TestMethod), nil)
 	replyChan := make(chan response, 1)
 	node.channel.enqueue(req, replyChan)
 
@@ -517,7 +518,7 @@ func TestChannelDeadlock(t *testing.T) {
 			ctx, cancel := context.WithTimeout(t.Context(), 3*time.Second)
 			defer cancel()
 
-			md := ordering.NewGorumsMetadata(ctx, uint64(100+id), mock.ServerMethodName)
+			md := ordering.NewGorumsMetadata(ctx, uint64(100+id), mock.TestMethod)
 			req := request{ctx: ctx, msg: NewRequestMessage(md, nil)}
 
 			// try to enqueue

config_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/relab/gorums"
 	"github.com/relab/gorums/internal/testutils/mock"
 	"google.golang.org/grpc/encoding"
+	pb "google.golang.org/protobuf/types/known/wrapperspb"
 )
 
 func init() {
@@ -226,8 +227,8 @@ func TestConfigConcurrentAccess(t *testing.T) {
 	for range 2 {
 		wg.Go(func() {
 			_, err := node.RPCCall(context.Background(), gorums.CallData{
-				Message: mock.NewRequest(""),
-				Method:  mock.ServerMethodName,
+				Message: pb.String(""),
+				Method:  mock.TestMethod,
 			})
 			if err != nil {
 				errCh <- err

internal/testutils/mock/mock.go

@@ -1,187 +1,133 @@
 package mock
 
 import (
-	"errors"
 	"fmt"
-	"sync"
-	"testing"
+	"strings"
 
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protodesc"
-	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/reflect/protoregistry"
 	"google.golang.org/protobuf/types/descriptorpb"
-	"google.golang.org/protobuf/types/dynamicpb"
+	"google.golang.org/protobuf/types/known/wrapperspb"
 )
 
-// ServerMethodName is the method supported by the mock package.
-const ServerMethodName = "mock.Server.Test"
-
-var (
-	// Mock Service Descriptors
-	mockFile     *descriptorpb.FileDescriptorProto
-	requestType  protoreflect.MessageType
-	responseType protoreflect.MessageType
+// TestMethod and GetValueMethod are the methods supported by the mock package.
+const (
+	TestMethod     = "mock.MockService.Test"
+	GetValueMethod = "mock.MockService.GetValue"
 )
 
-func init() {
-	// Initialize Mock Definitions
-	mockFile = &descriptorpb.FileDescriptorProto{
-		Name:    proto.String("mock/mock.proto"),
-		Package: proto.String("mock"),
-		MessageType: []*descriptorpb.DescriptorProto{
-			{
-				Name: proto.String("Request"),
-				Field: []*descriptorpb.FieldDescriptorProto{
-					{
-						Name:   proto.String("val"),
-						Number: proto.Int32(1),
-						Label:  descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL.Enum(),
-						Type:   descriptorpb.FieldDescriptorProto_TYPE_STRING.Enum(),
-					},
-				},
-			},
-			{
-				Name: proto.String("Response"),
-				Field: []*descriptorpb.FieldDescriptorProto{
-					{
-						Name:   proto.String("val"),
-						Number: proto.Int32(1),
-						Label:  descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL.Enum(),
-						Type:   descriptorpb.FieldDescriptorProto_TYPE_STRING.Enum(),
-					},
-				},
-			},
-		},
-		Service: []*descriptorpb.ServiceDescriptorProto{
-			{
-				Name: proto.String("Server"),
-				Method: []*descriptorpb.MethodDescriptorProto{
-					{
-						Name:       proto.String("Test"),
-						InputType:  proto.String(".mock.Request"),
-						OutputType: proto.String(".mock.Response"),
-					},
-				},
-			},
-		},
-	}
+// Service represents a service to be registered.
+type Service struct {
+	Name    string // Full package and service name, e.g., "mock.MockService"
+	Methods []Method
 }
 
-// Register registers the mock types in the global registry.
-// It is safe to call multiple times.
-func Register(t testing.TB) {
-	t.Helper()
-	if err := registerOnce(); err != nil {
-		t.Fatal(err)
-	}
+// Method represents a method in a service.
+type Method struct {
+	Name   string
+	Input  proto.Message
+	Output proto.Message
 }
 
-var registerOnce = sync.OnceValue(func() error {
-	if fd, err := protoregistry.GlobalFiles.FindFileByPath(mockFile.GetName()); err == nil {
-		// Already registered
-		return initDescriptors(fd)
-	}
+// RegisterServices registers the given services in the global registry.
+// It is safe to call multiple times, but services with the same package name
+// must be registered in the same call or be identical to previous registrations.
+// Returns an error if registration fails.
+func RegisterServices(services []Service) error {
+	// Group by package
+	packages := make(map[string][]*descriptorpb.ServiceDescriptorProto)
 
-	fd, err := protodesc.NewFile(mockFile, nil)
-	if err != nil {
-		return fmt.Errorf("failed to create mock file descriptor: %v", err)
-	}
-	if err := protoregistry.GlobalFiles.RegisterFile(fd); err != nil {
-		return fmt.Errorf("failed to register mock file: %v", err)
-	}
-	if err := initDescriptors(fd); err != nil {
-		return fmt.Errorf("failed to initialize mock descriptors: %v", err)
-	}
-	if err := protoregistry.GlobalTypes.RegisterMessage(requestType); err != nil {
-		return fmt.Errorf("failed to register Request type: %v", err)
-	}
-	if err := protoregistry.GlobalTypes.RegisterMessage(responseType); err != nil {
-		return fmt.Errorf("failed to register Response type: %v", err)
-	}
-	return nil
-})
+	for _, s := range services {
+		pkgName, svcName, found := strings.Cut(s.Name, ".")
+		if !found {
+			return fmt.Errorf("service name %q must contain a package", s.Name)
+		}
 
-func initDescriptors(fd protoreflect.FileDescriptor) error {
-	mockServiceDesc := fd.Services().ByName("Server")
-	if mockServiceDesc == nil {
-		return errors.New("service Server not found")
-	}
-	mockMethodDesc := mockServiceDesc.Methods().ByName("Test")
-	if mockMethodDesc == nil {
-		return errors.New("method Test not found")
-	}
-	requestMsgDesc := fd.Messages().ByName("Request")
-	if requestMsgDesc == nil {
-		return errors.New("message Request not found")
-	}
-	responseMsgDesc := fd.Messages().ByName("Response")
-	if responseMsgDesc == nil {
-		return errors.New("message Response not found")
+		svcDesc := &descriptorpb.ServiceDescriptorProto{
+			Name: proto.String(svcName),
+		}
+
+		for _, m := range s.Methods {
+			inDesc := m.Input.ProtoReflect().Descriptor()
+			outDesc := m.Output.ProtoReflect().Descriptor()
+			inName := string(inDesc.FullName())
+			outName := string(outDesc.FullName())
+
+			svcDesc.Method = append(svcDesc.Method, &descriptorpb.MethodDescriptorProto{
+				Name:       proto.String(m.Name),
+				InputType:  proto.String("." + inName),
+				OutputType: proto.String("." + outName),
+			})
+		}
+		packages[pkgName] = append(packages[pkgName], svcDesc)
 	}
-	requestType = dynamicpb.NewMessageType(requestMsgDesc)
-	responseType = dynamicpb.NewMessageType(responseMsgDesc)
-	return nil
-}
 
-// Helpers for Mock messages
+	for pkg, svcDescriptors := range packages {
+		// Collect dependencies
+		deps := make(map[string]struct{})
 
-const panicMsg = "mock.Register() must be called before using NewRequest/NewResponse"
+		// Iterate over the original services to find dependencies for this package.
+		for _, s := range services {
+			pName, _, found := strings.Cut(s.Name, ".")
+			if !found || pName != pkg {
+				continue
+			}
 
-func NewRequest(val string) proto.Message {
-	if requestType == nil {
-		panic(panicMsg)
-	}
-	msg := requestType.New()
-	if val != "" {
-		fd := msg.Descriptor().Fields().ByName("val")
-		if fd != nil {
-			msg.Set(fd, protoreflect.ValueOfString(val))
+			for _, m := range s.Methods {
+				if d := m.Input.ProtoReflect().Descriptor().ParentFile(); d != nil {
+					deps[d.Path()] = struct{}{}
+				}
+				if d := m.Output.ProtoReflect().Descriptor().ParentFile(); d != nil {
+					deps[d.Path()] = struct{}{}
+				}
+			}
 		}
-	}
-	return msg.Interface()
-}
 
-func NewResponse(val string) proto.Message {
-	if responseType == nil {
-		panic(panicMsg)
-	}
-	msg := responseType.New()
-	if val != "" {
-		fd := msg.Descriptor().Fields().ByName("val")
-		if fd != nil {
-			msg.Set(fd, protoreflect.ValueOfString(val))
+		fd := &descriptorpb.FileDescriptorProto{
+			Name:    proto.String(fmt.Sprintf("mock/%s.proto", pkg)),
+			Package: proto.String(pkg),
+			Service: svcDescriptors,
+		}
+
+		for dep := range deps {
+			fd.Dependency = append(fd.Dependency, dep)
+		}
+
+		// Check if already registered
+		if _, err := protoregistry.GlobalFiles.FindFileByPath(fd.GetName()); err == nil {
+			continue // Already registered
+		}
+
+		fileDesc, err := protodesc.NewFile(fd, protoregistry.GlobalFiles)
+		if err != nil {
+			return fmt.Errorf("failed to create file descriptor for %s: %v", pkg, err)
+		}
+
+		if err := protoregistry.GlobalFiles.RegisterFile(fileDesc); err != nil {
+			return fmt.Errorf("failed to register file %s: %v", pkg, err)
 		}
 	}
-	return msg.Interface()
+	return nil
 }
 
+// Helpers for Mock messages
+
 func GetVal(msg proto.Message) string {
 	if msg == nil {
 		return ""
 	}
-	m := msg.ProtoReflect()
-	if !m.IsValid() {
-		return ""
-	}
-	fd := m.Descriptor().Fields().ByName("val")
-	if fd == nil {
-		return ""
+	if m, ok := msg.(*wrapperspb.StringValue); ok {
+		return m.Value
 	}
-	return m.Get(fd).String()
+	return ""
 }
 
 func SetVal(msg proto.Message, val string) {
 	if msg == nil {
 		return
 	}
-	m := msg.ProtoReflect()
-	if !m.IsValid() {
-		return
-	}
-	fd := m.Descriptor().Fields().ByName("val")
-	if fd == nil {
-		return
+	if m, ok := msg.(*wrapperspb.StringValue); ok {
+		m.Value = val
 	}
-	m.Set(fd, protoreflect.ValueOfString(val))
 }

main_test.go

@@ -0,0 +1,34 @@
+package gorums
+
+import (
+	"os"
+	"testing"
+
+	"github.com/relab/gorums/internal/testutils/mock"
+	pb "google.golang.org/protobuf/types/known/wrapperspb"
+)
+
+func TestMain(m *testing.M) {
+	// Register the default mock services for integration tests.
+	if err := mock.RegisterServices([]mock.Service{
+		{
+			Name: "mock.MockService",
+			Methods: []mock.Method{
+				{
+					Name:   "Test",
+					Input:  &pb.StringValue{},
+					Output: &pb.StringValue{},
+				},
+				{
+					Name:   "GetValue",
+					Input:  &pb.Int32Value{},
+					Output: &pb.Int32Value{},
+				},
+			},
+		},
+	}); err != nil {
+		panic(err)
+	}
+
+	os.Exit(m.Run())
+}