some tests added

cmd/protoc-gen-python-grpc/internal/flags/root.go

@@ -1,10 +1,13 @@
 package flags
 
 import (
-	"github.com/pseudomuto/protokit"
+	"os"
+
 	"github.com/spf13/cobra"
 
 	"github.com/Djarvur/protoc-gen-python-grpc/internal/generator"
+	"github.com/Djarvur/protoc-gen-python-grpc/internal/kit"
+	"github.com/Djarvur/protoc-gen-python-grpc/internal/template"
 )
 
 const (
@@ -14,7 +17,7 @@ const (
 )
 
 func Root() *cobra.Command {
-	templateSource := newTemplateValue()
+	templateSource := template.NewTemplateValue()
 
 	cmd := &cobra.Command{ //nolint:exhaustruct
 		Use:   "protoc-gen-python-grpc",
@@ -23,7 +26,7 @@ func Root() *cobra.Command {
 		Run: func(cmd *cobra.Command, args []string) {
 			runRoot(
 				must(cmd.Flags().GetString(fileNameSuffixFlag)),
-				templateSource.source,
+				templateSource.Source(),
 			)
 		},
 	}
@@ -35,7 +38,7 @@ func Root() *cobra.Command {
 }
 
 func runRoot(suffix, templateSource string) {
-	if err := protokit.RunPlugin(must(generator.New(suffix, templateSource))); err != nil {
+	if err := kit.New().RunPluginWithIO(generator.New(suffix, templateSource), os.Stdin, os.Stdout); err != nil {
 		panic(err)
 	}
 }

internal/generator/generator.go

@@ -43,24 +43,24 @@ var _ protokit.Plugin = (*generator)(nil)
 // It's an implementation of generator from github.com/pseudomuto/protokit.
 type generator struct {
 	Suffix   string
-	Template *template.Template
+	Template string
 }
 
-func New(suffix, tmplSrc string) (*generator, error) {
-	tmpl, err := buildTemplate(tmplSrc)
-	if err != nil {
-		return nil, err
-	}
-
+func New(suffix, tmplSrc string) *generator {
 	return &generator{
 		Suffix:   suffix,
-		Template: tmpl,
-	}, nil
+		Template: tmplSrc,
+	}
 }
 
 // Generate compiles the documentation and generates the CodeGeneratorResponse to send back to protoc. It does this
 // by rendering a template based on the options parsed from the CodeGeneratorRequest.
 func (p *generator) Generate(r *pluginpb.CodeGeneratorRequest) (*pluginpb.CodeGeneratorResponse, error) {
+	tmpl, err := buildTemplate(p.Template)
+	if err != nil {
+		return nil, fmt.Errorf("building template: %w", err)
+	}
+
 	resp := new(pluginpb.CodeGeneratorResponse)
 
 	for _, fds := range protokit.ParseCodeGenRequest(r) {
@@ -70,7 +70,7 @@ func (p *generator) Generate(r *pluginpb.CodeGeneratorRequest) (*pluginpb.CodeGe
 			Services: buildServices(fds.GetServices()),
 		}
 
-		content, errExecute := executeTemplate(p.Template, data)
+		content, errExecute := executeTemplate(tmpl, data)
 		if errExecute != nil {
 			return nil, errExecute
 		}
@@ -151,3 +151,11 @@ func buildMethods(in []*protokit.MethodDescriptor) []Method {
 
 	return out
 }
+
+func must[T any](v T, err error) T {
+	if err != nil {
+		panic(err)
+	}
+
+	return v
+}

internal/kit/kit.go

@@ -0,0 +1,46 @@
+package kit
+
+import (
+	"bytes"
+	"io"
+	"os"
+
+	plugin_go "github.com/golang/protobuf/protoc-gen-go/plugin"
+	"github.com/pseudomuto/protokit"
+)
+
+type Plugin interface {
+	Generate(req *plugin_go.CodeGeneratorRequest) (*plugin_go.CodeGeneratorResponse, error)
+}
+
+type Kit struct {
+}
+
+func New() Kit {
+	return Kit{}
+}
+
+func (k Kit) RunPluginWithIO(p Plugin, r io.Reader, w io.Writer) error {
+	in, err := io.ReadAll(r)
+	if err != nil {
+		panic(err)
+	}
+
+	if err = os.WriteFile("testdata/in.bin", in, 0o644); err != nil {
+		panic(err)
+	}
+
+	var out bytes.Buffer
+
+	errPlugin := protokit.RunPluginWithIO(p, bytes.NewBuffer(in), &out)
+
+	if err = os.WriteFile("testdata/out.bin", out.Bytes(), 0o644); err != nil {
+		panic(err)
+	}
+
+	if _, err = io.Copy(w, &out); err != nil {
+		panic(err)
+	}
+
+	return errPlugin
+}

internal/kit/kit_test.go

@@ -0,0 +1,39 @@
+package kit_test
+
+import (
+	"bytes"
+	_ "embed"
+	"testing"
+
+	"github.com/Djarvur/protoc-gen-python-grpc/internal/generator"
+	"github.com/Djarvur/protoc-gen-python-grpc/internal/kit"
+	"github.com/Djarvur/protoc-gen-python-grpc/internal/template"
+	"github.com/stretchr/testify/require"
+)
+
+const (
+	fileNameSuffixDefault = "_pb2_grpc.py"
+)
+
+var (
+	//go:embed testdata/in.bin
+	inBytes []byte
+
+	//go:embed testdata/out.bin
+	outBytes []byte
+)
+
+func TestRunPluginWithIO(t *testing.T) {
+	t.Parallel()
+
+	out := &bytes.Buffer{}
+
+	err := kit.New().RunPluginWithIO(
+		generator.New(fileNameSuffixDefault, template.NewTemplateValue().Source()),
+		bytes.NewBuffer(inBytes),
+		out,
+	)
+	require.NoError(t, err)
+
+	require.Equal(t, out.Bytes(), outBytes)
+}

internal/kit/testdata/out.bin

@@ -0,0 +1,67 @@
+
z�
+test/v1/calls_pb2_grpc.pyz�# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+
+from test.v1 import calls_pb2 as test_dot_v1_dot_calls__pb2
+
+
+class CallsServiceStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.Call = channel.unary_unary(
+                '/test.v1.CallsService/Call',
+                request_serializer=test_dot_v1_dot_calls__pb2.CallRequest.SerializeToString,
+                response_deserializer=test_dot_v1_dot_calls__pb2.CallResponse.FromString,
+                )
+
+
+class CallsServiceServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def Call(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_CallsServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'Call': grpc.unary_unary_rpc_method_handler(
+                    servicer.Call,
+                    request_deserializer=test_dot_v1_dot_calls__pb2.CallRequest.FromString,
+                    response_serializer=test_dot_v1_dot_calls__pb2.CallResponse.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'test.v1.CallsService', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+
+
+ # This class is part of an EXPERIMENTAL API.
+class CallsService(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def Call(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/test.v1.CallsService/Call',
+            test_dot_v1_dot_calls__pb2.CallRequest.SerializeToString,
+            test_dot_v1_dot_calls__pb2.CallResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)

cmd/protoc-gen-python-grpc/internal/flags/source.go → internal/template/template.go

@@ -1,4 +1,4 @@
-package flags
+package template
 
 import (
 	_ "embed"
@@ -11,26 +11,26 @@ import (
 //go:embed pb2_grpc.py.tmpl
 var defaultTemplateSrc string
 
-var _ pflag.Value = (*sourceValue)(nil)
+var _ pflag.Value = (*TemplateValue)(nil)
 
-type sourceValue struct {
+type TemplateValue struct {
 	name   string
 	source string
 }
 
-func (r *sourceValue) String() string {
+func (r *TemplateValue) String() string {
 	return r.name
 }
 
-func newTemplateValue() *sourceValue {
-	return &sourceValue{
+func NewTemplateValue() *TemplateValue {
+	return &TemplateValue{
 		name:   "EMBEDDED",
 		source: defaultTemplateSrc,
 	}
 }
 
 // Set is a method to set the template value.
-func (r *sourceValue) Set(s string) error {
+func (r *TemplateValue) Set(s string) error {
 	b, err := os.ReadFile(s)
 	if err != nil {
 		return fmt.Errorf("reading template %q: %w", s, err)
@@ -42,6 +42,10 @@ func (r *sourceValue) Set(s string) error {
 }
 
 // Type required to implement pflag.Value.
-func (*sourceValue) Type() string {
+func (*TemplateValue) Type() string {
 	return "text/template"
 }
+
+func (v *TemplateValue) Source() string {
+	return v.source
+}