mock.go

package main

import (
	"strings"

	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/types/descriptorpb"
)

const deprecationComment = "// Deprecated: Do not use."

const (
	contextPackage = protogen.GoImportPath("context")
	netPackage     = protogen.GoImportPath("net")
	grpcPackage    = protogen.GoImportPath("google.golang.org/grpc")
	mockPackage    = protogen.GoImportPath("github.com/stretchr/testify/mock")
)

func generateFile(gen *protogen.Plugin, file *protogen.File) *protogen.GeneratedFile {
	if len(file.Services) == 0 {
		return nil
	}

	filename := file.GeneratedFilenamePrefix + "_grpc_mock.pb.go"
	g := gen.NewGeneratedFile(filename, file.GoImportPath)

	generateFileHeader(gen, file, g)
	generateFileContent(gen, file, g)

	return g
}

func generateFileHeader(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile) {
	g.P("// Code generated by protoc-gen-go-grpc-mock (https://github.com/nefixestrada/protoc-gen-go-grpc-mock). DO NOT EDIT.")
	g.P("// version: ", version)
	g.P()
	g.P("package ", file.GoPackageName)
	g.P()
}

func generateFileContent(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile) {
	if len(file.Services) == 0 {
		return
	}

	g.P("// This is a compile-time assertion to ensure that this generated file")
	g.P("// is compatible with the grpc package it is being compiled against.")
	g.P("const _ = ", grpcPackage.Ident("SupportPackageIsVersion7"))
	g.P()

	for _, service := range file.Services {
		genServiceMock(gen, file, g, service)
	}
}

func genServiceMock(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile, service *protogen.Service) {
	serviceType := service.GoName + "ServiceMock"

	g.P("// New", serviceType, " creates and starts a ", service.GoName, " service mock.")
	g.P("// it returns the mock and a client for the ")
	g.P("func New", serviceType, "() (*", serviceType, ", ", service.GoName, "Client, error) {")
	g.P("addr, err := ", g.QualifiedGoIdent(netPackage.Ident("ResolveTCPAddr")), "(\"tcp\", \"localhost:0\")")
	g.P("if err != nil {")
	g.P("return nil, nil, err")
	g.P("}")
	g.P()
	g.P("lis, err := ", g.QualifiedGoIdent(netPackage.Ident("ListenTCP")), "(\"tcp\", addr)")
	g.P("if err != nil {")
	g.P("return nil, nil, err")
	g.P("}")
	g.P()
	g.P("s := ", g.QualifiedGoIdent(grpcPackage.Ident("NewServer")), "()")
	g.P("m := &", serviceType, "{server: s}")
	g.P()
	g.P("Register", service.GoName+"Server(s, m)")
	g.P()
	g.P("go s.Serve(lis)")
	g.P()
	g.P("cc, err := ", g.QualifiedGoIdent(grpcPackage.Ident("Dial")), "(lis.Addr().String(), ", g.QualifiedGoIdent(grpcPackage.Ident("WithInsecure")), "())")
	g.P("if err != nil {")
	g.P("return nil, nil, err")
	g.P("}")
	g.P("m.cc = cc")
	g.P()
	g.P("cli := New", service.GoName, "Client(cc)")
	g.P()
	g.P("return m, cli, nil")
	g.P("}")
	g.P()

	g.P("// ", serviceType, " is a mock for the ", service.GoName, "Service service.")
	if service.Desc.Options().(*descriptorpb.ServiceOptions).GetDeprecated() {
		g.P("//")
		g.P(deprecationComment)
	}
	g.Annotate(serviceType, service.Location)
	g.P("type ", serviceType, " struct {")
	g.P(g.QualifiedGoIdent(mockPackage.Ident("Mock")))
	g.P("server *", g.QualifiedGoIdent(grpcPackage.Ident("Server")))
	g.P("cc *", g.QualifiedGoIdent(grpcPackage.Ident("ClientConn")))
	g.P("Unimplemented", service.GoName, "Server")
	g.P("}")
	g.P()

	g.P("func (m *", serviceType, ") Stop() {")
	g.P("m.cc.Close()")
	g.P("m.server.Stop()")
	g.P("}")
	g.P()

	for _, method := range service.Methods {
		genMethodHandler(gen, g, method)
	}
}

func handlerSignature(g *protogen.GeneratedFile, method *protogen.Method) string {
	var reqArgs []string
	ret := "error"
	if !method.Desc.IsStreamingClient() && !method.Desc.IsStreamingServer() {
		reqArgs = append(reqArgs, "ctx "+g.QualifiedGoIdent(contextPackage.Ident("Context")))
		ret = "(*" + g.QualifiedGoIdent(method.Output.GoIdent) + ", error)"
	}
	if !method.Desc.IsStreamingClient() {
		reqArgs = append(reqArgs, "in *"+g.QualifiedGoIdent(method.Input.GoIdent))
	}
	if method.Desc.IsStreamingClient() || method.Desc.IsStreamingServer() {
		reqArgs = append(reqArgs, method.Parent.GoName+"_"+method.GoName+"ServerMock")
	}
	return "(" + strings.Join(reqArgs, ", ") + ") " + ret
}

func genMethodHandler(gen *protogen.Plugin, g *protogen.GeneratedFile, method *protogen.Method) {
	service := method.Parent

	g.P("const ", method.GoName, "= \"", method.GoName, "\"")

	g.P("func (m *", service.GoName, "ServiceMock) ", method.GoName, handlerSignature(g, method), " {")

	// Unary
	if !method.Desc.IsStreamingClient() && !method.Desc.IsStreamingServer() {
		g.P("args := m.Called(ctx, in)")
		g.P("return args.Get(0).(*", method.Output.GoIdent, "), args.Error(1)")
	}

	g.P("}")
	g.P()
}