main.go

package main

import (
	"flag"
	"fmt"
	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/types/descriptorpb"
	"strconv"
	"strings"
)

const (
	fmtPackage         = protogen.GoImportPath("fmt")
	contextPackage     = protogen.GoImportPath("context")
	protoPackage       = protogen.GoImportPath("github.com/golang/protobuf/proto")
	grpcPackage        = protogen.GoImportPath("google.golang.org/grpc")
	grpchandlerPackage = protogen.GoImportPath("github.com/carvalhorr/protoc-gen-mock/grpchandler")
	stubPackage        = protogen.GoImportPath("github.com/carvalhorr/protoc-gen-mock/stub")
	deprecationComment = "// Deprecated: Do not use."
)

func main() {
	var (
		flags flag.FlagSet
		//plugins      = flags.String("plugins", "", "list of plugins to enable (supported values: grpc)")
		importPrefix = flags.String("import_prefix", "", "prefix to prepend to import paths")
	)
	importRewriteFunc := func(importPath protogen.GoImportPath) protogen.GoImportPath {
		switch importPath {
		case "context", "fmt", "math":
			return importPath
		}
		if *importPrefix != "" {
			return protogen.GoImportPath(*importPrefix) + importPath
		}
		return importPath
	}
	protogen.Options{
		ParamFunc:         flags.Set,
		ImportRewriteFunc: importRewriteFunc,
	}.Run(func(gen *protogen.Plugin) error {
		for _, f := range gen.Files {
			GenerateFile(gen, f)
		}
		return nil
	})
}

// GenerateFile generates a _grpc.pb.go file containing gRPC service definitions.
func GenerateFile(gen *protogen.Plugin, file *protogen.File) *protogen.GeneratedFile {
	if len(file.Services) == 0 {
		return nil
	}
	filename := file.GeneratedFilenamePrefix + ".mock.pb.go"
	g := gen.NewGeneratedFile(filename, file.GoImportPath)
	mockGenerator := mockServicesGenerator{
		gen:  gen,
		file: file,
		g:    g,
	}
	mockGenerator.genHeader(string(file.GoPackageName))
	mockGenerator.GenerateFileContent()
	return g
}

type mockServicesGenerator struct {
	gen  *protogen.Plugin
	file *protogen.File
	g    *protogen.GeneratedFile
}

// GenerateFileContent generates the gRPC service definitions, excluding the package statement.
func (m mockServicesGenerator) GenerateFileContent() {
	if len(m.file.Services) == 0 {
		return
	}
	for _, service := range m.file.Services {
		m.genService(service)
	}
}

func (m mockServicesGenerator) genService(service *protogen.Service) {
	m.genMockServiceConstructor(service)
	m.genMockServiceDefinition(service)
	m.genMockServiceRegistrationFunction(service)
	m.genGetSupportedMethodsFunction(service)
	m.genGetPayloadExamplesFunction(service)
	m.genGetRequestInstance(service)
	m.genGetResponseInstance(service)
	m.genGetStubsValidator(service)
	m.genIsValid(service)
	m.genForwardRequest(service)
	m.genMockServiceDescriptor(service)
	for _, method := range service.Methods {
		methodHandlerName := m.getMethodHandlerName(service, method)
		m.genMockMethodHandler(service, method, methodHandlerName)
	}
}

func (m mockServicesGenerator) genHeader(packageName string) {
	m.g.P("// Code generated by protoc-gen-mock. DO NOT EDIT.")
	m.g.P()
	m.g.P("package ", packageName)
	m.g.P()
}

func (m mockServicesGenerator) genMockServiceConstructor(service *protogen.Service) {
	serviceName := m.getMockServiceName(service)
	m.g.P("func New", serviceName, "(stubsMatcher ", stubPackage.Ident("StubsMatcher"), ") ", grpchandlerPackage.Ident("MockService"), "{")
	m.g.P("return &", unexport(serviceName), "{")
	m.g.P("StubsMatcher: stubsMatcher,")
	m.g.P("}")
	m.g.P("}")
}

func (m mockServicesGenerator) genMockServiceDefinition(service *protogen.Service) {
	serviceName := m.getMockServiceName(service)
	serviceBaseInterfaceName := m.getMockServerBaseInterfaceName(service)
	m.g.P("// ", serviceName, " is the server API for ", service.GoName, " service.")
	if service.Desc.Options().(*descriptorpb.ServiceOptions).GetDeprecated() {
		m.g.P("//")
		m.g.P(deprecationComment)
	}
	m.g.Annotate(serviceName, service.Location)
	m.g.P("type ", unexport(serviceName), " struct {")
	m.g.P(serviceBaseInterfaceName)
	m.g.P("StubsMatcher ", stubPackage.Ident("StubsMatcher"))
	m.g.P("}")
	m.g.P()
}

func (m mockServicesGenerator) genMockServiceRegistrationFunction(service *protogen.Service) {
	// Server registration.
	if service.Desc.Options().(*descriptorpb.ServiceOptions).GetDeprecated() {
		m.g.P(deprecationComment)
	}
	m.g.P("func (mock *", unexport(m.getMockServiceName(service)), ") Register(s *", grpcPackage.Ident("Server"), ") {")
	m.g.P("s.RegisterService(&", m.getMockServiceDescriptorName(service), ", ", m.getMockServerBaseInterfaceName(service), "(mock))")
	m.g.P("}")
	m.g.P()
}

func (m mockServicesGenerator) genGetSupportedMethodsFunction(service *protogen.Service) {
	m.g.P("func (mock *", unexport(m.getMockServiceName(service)), ") GetSupportedMethods() []string {")
	m.g.P("return []string{")
	for _, method := range service.Methods {
		m.g.P(strconv.Quote(fmt.Sprintf("/%s/%s", service.Desc.FullName(), method.GoName)), ",")
	}
	m.g.P("}")
	m.g.P("}")
	m.g.P()
}

func (m mockServicesGenerator) genGetPayloadExamplesFunction(service *protogen.Service) {
	m.g.P("func (mock *", unexport(m.getMockServiceName(service)), ") GetPayloadExamples() []", stubPackage.Ident("Stub"), "{")
	m.g.P("return []", stubPackage.Ident("Stub"), "{")
	for _, method := range service.Methods {
		m.g.P("{")
		m.g.P("FullMethod: ", strconv.Quote(fmt.Sprintf("/%s/%s", service.Desc.FullName(), method.GoName)), ",")
		m.g.P("Type: ", strconv.Quote("mock | forward"), ",")
		m.g.P("Request: &", stubPackage.Ident("StubRequest"), " {")
		m.g.P("Match: \"exact | partial\",")
		m.g.P("Content: ", stubPackage.Ident("JsonString"), "(", stubPackage.Ident("CreateStubExample"), "(new(", method.Input.GoIdent, "))", "),")
		m.g.P("Metadata: make(map[string][]string, 0),")
		m.g.P("},")
		m.g.P("Response: &", stubPackage.Ident("StubResponse"), " {")
		m.g.P("Type: ", strconv.Quote("success | error"), ", ")
		m.g.P("Content: ", stubPackage.Ident("JsonString"), "(", stubPackage.Ident("CreateStubExample"), "(new(", method.Output.GoIdent, "))", "),")
		m.g.P("},")
		m.g.P("Forward: &", stubPackage.Ident("StubForward"), " {")
		m.g.P("ServerAddress: ", strconv.Quote("yourserver:port"), ",")
		m.g.P("Record: true,")
		m.g.P("},")
		m.g.P("},")
	}
	m.g.P("}")
	m.g.P("}")
	m.g.P("")
}

func (m mockServicesGenerator) genGetRequestInstance(service *protogen.Service) {
	m.g.P("func (mock *", unexport(m.getMockServiceName(service)), ") GetRequestInstance(methodName string) ", protoPackage.Ident("Message"), " {")
	m.g.P("switch methodName {")
	for _, method := range service.Methods {
		m.g.P("case ", strconv.Quote(fmt.Sprintf("/%s/%s", service.Desc.FullName(), method.GoName)), ":")
		m.g.P("return new(", method.Input.GoIdent, ")")
	}
	m.g.P("}")
	m.g.P("return nil")
	m.g.P("}")
	m.g.P("")
}

func (m mockServicesGenerator) genGetResponseInstance(service *protogen.Service) {
	m.g.P("func (mock *", unexport(m.getMockServiceName(service)), ") GetResponseInstance(methodName string) ", protoPackage.Ident("Message"), "{")
	m.g.P("switch methodName {")
	for _, method := range service.Methods {
		m.g.P("case ", strconv.Quote(fmt.Sprintf("/%s/%s", service.Desc.FullName(), method.GoName)), ":")
		m.g.P("return new(", method.Output.GoIdent, ")")
	}
	m.g.P("}")
	m.g.P("return nil")
	m.g.P("}")
	m.g.P("")
}
func (m mockServicesGenerator) genMockServiceDescriptor(service *protogen.Service) {
	// Service descriptor.
	m.g.P("var ", m.getMockServiceDescriptorName(service), " = ", grpcPackage.Ident("ServiceDesc"), " {")
	m.g.P("ServiceName: ", strconv.Quote(string(service.Desc.FullName())), ",")
	m.g.P("HandlerType: (*", m.getMockServerBaseInterfaceName(service), ")(nil),")
	m.g.P("Methods: []", grpcPackage.Ident("MethodDesc"), "{")
	for _, method := range service.Methods {
		methodHandlerName := m.getMethodHandlerName(service, method)
		if method.Desc.IsStreamingClient() || method.Desc.IsStreamingServer() {
			continue // skip if it is streaming
		}
		m.g.P("{")
		m.g.P("MethodName: ", strconv.Quote(string(method.Desc.Name())), ",")
		m.g.P("Handler: ", methodHandlerName, ",")
		m.g.P("},")
	}
	m.g.P("},")
	m.g.P("Streams: []", grpcPackage.Ident("StreamDesc"), "{")
	for _, method := range service.Methods {
		methodHandlerName := m.getMethodHandlerName(service, method)
		if !method.Desc.IsStreamingClient() && !method.Desc.IsStreamingServer() {
			continue
		}
		m.g.P("{")
		m.g.P("StreamName: ", strconv.Quote(string(method.Desc.Name())), ",")
		m.g.P("Handler: ", methodHandlerName, ",")
		if method.Desc.IsStreamingServer() {
			m.g.P("ServerStreams: true,")
		}
		if method.Desc.IsStreamingClient() {
			m.g.P("ClientStreams: true,")
		}
		m.g.P("},")
	}
	m.g.P("},")
	m.g.P("Metadata: \"", m.file.Desc.Path(), "\",")
	m.g.P("}")
	m.g.P()
}

func (m mockServicesGenerator) genMockMethodHandler(service *protogen.Service, method *protogen.Method, hname string) {
	if !method.Desc.IsStreamingClient() && !method.Desc.IsStreamingServer() {
		m.g.P("func ", hname, "(srv interface{}, ctx ", contextPackage.Ident("Context"), ", dec func(interface{}) error, interceptor ", grpcPackage.Ident("UnaryServerInterceptor"), ") (interface{}, error) {")
		m.g.P("in := new(", method.Input.GoIdent, ")")
		m.g.P("if err := dec(in); err != nil { return nil, err }")
		m.g.P("out := new(", method.Output.GoIdent, ")")
		m.g.P("fullMethod := ", strconv.Quote(fmt.Sprintf("/%s/%s", service.Desc.FullName(), method.GoName)))
		m.g.P("stubsMatcher := (srv).(*", unexport(m.getMockServiceName(service)), ").StubsMatcher")
		m.g.P("return ", grpchandlerPackage.Ident("MockHandler"), "(ctx, stubsMatcher, fullMethod, in, out)")
		m.g.P("}")
		m.g.P()
		return
	}
	m.g.P("func ", hname, "(srv interface{}, stream ", grpcPackage.Ident("ServerStream"), ") error {")
	m.g.P("// Mock not implemented for streaming")
	m.g.P("return ", fmtPackage.Ident("Errorf"), "(\"mock not implemented for streaming\")")
	m.g.P("}")
	m.g.P()
}

func (m mockServicesGenerator) genGetStubsValidator(service *protogen.Service) {
	m.g.P("func(mock *", unexport(m.getMockServiceName(service)), ") GetStubsValidator() ", stubPackage.Ident("StubsValidator"), "{")
	m.g.P("return mock")
	m.g.P("}")
	m.g.P("")
}

func (m mockServicesGenerator) genIsValid(service *protogen.Service) {

	m.g.P("func(mock *", unexport(m.getMockServiceName(service)), ") IsValid(s *", stubPackage.Ident("Stub"), ") (isValid bool, errorMessages []string) {")
	m.g.P("switch s.FullMethod {")
	for _, method := range service.Methods {
		m.g.P("case ", strconv.Quote(fmt.Sprintf("/%s/%s", service.Desc.FullName(), method.GoName)), ":")
		m.g.P("req := new(", method.Input.GoIdent, ")")
		m.g.P("resp := new(", method.Output.GoIdent, ")")
		m.g.P("return ", stubPackage.Ident("IsStubValid"), "(s, req.ProtoReflect().Descriptor(), resp.ProtoReflect().Descriptor())")
	}
	m.g.P("default:")
	m.g.P("return true, nil")
	m.g.P("}")
	m.g.P("}")
	m.g.P("")

}

func (m mockServicesGenerator) genForwardRequest(service *protogen.Service) {

	m.g.P("func(mock *", unexport(m.getMockServiceName(service)), ") ForwardRequest(conn grpc.ClientConnInterface, ctx context.Context, methodName string, req interface{}) (interface{}, error) {")
	m.g.P("client := New", service.Desc.Name(), "Client(conn)")
	m.g.P("switch methodName {")
	for _, method := range service.Methods {
		m.g.P("case ", strconv.Quote(fmt.Sprintf("/%s/%s", service.Desc.FullName(), method.GoName)), ":")
		m.g.P("return client.", method.GoName, "(ctx, req.(*", method.Input.GoIdent, "))")
	}
	m.g.P("}")
	m.g.P("return nil, nil")
	m.g.P("}")
	m.g.P("")

}

func (m mockServicesGenerator) getMockServiceName(service *protogen.Service) string {
	return service.GoName + "MockService"
}

func (m mockServicesGenerator) getMockServerBaseInterfaceName(service *protogen.Service) string {
	return service.GoName + "Server"
}

func (m mockServicesGenerator) getMockServiceDescriptorName(service *protogen.Service) string {
	return "_" + service.GoName + "_MockServiceDesc"
}

func (m mockServicesGenerator) getMethodHandlerName(service *protogen.Service, method *protogen.Method) string {
	return fmt.Sprintf("_%s_%s_MockHandler", service.GoName, method.GoName)
}

func unexport(s string) string { return strings.ToLower(s[:1]) + s[1:] }