From df2799fff112c109f6ac5872af49b7c7a4ee9734 Mon Sep 17 00:00:00 2001
From: zhaolei <2290938350@qq.com>
Date: Sat, 4 Nov 2023 20:42:25 +0800
Subject: [PATCH] fix import error if generate multiple proto (#3694)

---
 tools/goctl/rpc/generator/genpb.go      | 14 +++++++++-----
 tools/goctl/rpc/generator/genpb_test.go | 20 ++++++++++----------
 2 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/tools/goctl/rpc/generator/genpb.go b/tools/goctl/rpc/generator/genpb.go
index 797a1fc8..758bc2f6 100644
--- a/tools/goctl/rpc/generator/genpb.go
+++ b/tools/goctl/rpc/generator/genpb.go
@@ -31,14 +31,14 @@ func (g *Generator) genPbDirect(ctx DirContext, c *ZRpcContext) error {
 }
 
 func (g *Generator) setPbDir(ctx DirContext, c *ZRpcContext) error {
-	pbDir, err := findPbFile(c.GoOutput, false)
+	pbDir, err := findPbFile(c.GoOutput, c.Src, false)
 	if err != nil {
 		return err
 	}
 	if len(pbDir) == 0 {
 		return fmt.Errorf("pg.go is not found under %q", c.GoOutput)
 	}
-	grpcDir, err := findPbFile(c.GrpcOutput, true)
+	grpcDir, err := findPbFile(c.GrpcOutput, c.Src, true)
 	if err != nil {
 		return err
 	}
@@ -62,7 +62,11 @@ const (
 	grpcSuffix = "_grpc.pb.go"
 )
 
-func findPbFile(current string, grpc bool) (string, error) {
+func findPbFile(current string, src string, grpc bool) (string, error) {
+	protoName := strings.TrimSuffix(filepath.Base(src), filepath.Ext(src))
+	pbFile := protoName + "." + pbSuffix
+	grpcFile := protoName + grpcSuffix
+
 	fileSystem := os.DirFS(current)
 	var ret string
 	err := fs.WalkDir(fileSystem, ".", func(path string, d fs.DirEntry, err error) error {
@@ -71,11 +75,11 @@ func findPbFile(current string, grpc bool) (string, error) {
 		}
 		if strings.HasSuffix(path, pbSuffix) {
 			if grpc {
-				if strings.HasSuffix(path, grpcSuffix) {
+				if strings.HasSuffix(path, grpcFile) {
 					ret = path
 					return os.ErrExist
 				}
-			} else if !strings.HasSuffix(path, grpcSuffix) {
+			} else if strings.HasSuffix(path, pbFile) {
 				ret = path
 				return os.ErrExist
 			}
diff --git a/tools/goctl/rpc/generator/genpb_test.go b/tools/goctl/rpc/generator/genpb_test.go
index 9ed939a4..dbc6d5d9 100644
--- a/tools/goctl/rpc/generator/genpb_test.go
+++ b/tools/goctl/rpc/generator/genpb_test.go
@@ -46,12 +46,12 @@ service Greeter {
 			t.Log(err)
 			return
 		}
-		pbDir, err := findPbFile(output, false)
+		pbDir, err := findPbFile(output, protoFile, false)
 		assert.Nil(t, err)
 		pbGo := filepath.Join(pbDir, "greet.pb.go")
 		assert.True(t, pathx.FileExists(pbGo))
 
-		grpcDir, err := findPbFile(output, true)
+		grpcDir, err := findPbFile(output, protoFile, true)
 		assert.Nil(t, err)
 		grpcGo := filepath.Join(grpcDir, "greet_grpc.pb.go")
 		assert.True(t, pathx.FileExists(grpcGo))
@@ -76,12 +76,12 @@ service Greeter {
 			t.Log(err)
 			return
 		}
-		pbDir, err := findPbFile(output, false)
+		pbDir, err := findPbFile(output, protoFile, false)
 		assert.Nil(t, err)
 		pbGo := filepath.Join(pbDir, "greet.pb.go")
 		assert.True(t, pathx.FileExists(pbGo))
 
-		grpcDir, err := findPbFile(output, true)
+		grpcDir, err := findPbFile(output, protoFile, true)
 		assert.Nil(t, err)
 		grpcGo := filepath.Join(grpcDir, "greet_grpc.pb.go")
 		assert.True(t, pathx.FileExists(grpcGo))
@@ -108,12 +108,12 @@ service Greeter {
 			t.Log(err)
 			return
 		}
-		pbDir, err := findPbFile(output, false)
+		pbDir, err := findPbFile(output, protoFile, false)
 		assert.Nil(t, err)
 		pbGo := filepath.Join(pbDir, "greet.pb.go")
 		assert.True(t, pathx.FileExists(pbGo))
 
-		grpcDir, err := findPbFile(output, true)
+		grpcDir, err := findPbFile(output, protoFile, true)
 		assert.Nil(t, err)
 		grpcGo := filepath.Join(grpcDir, "greet_grpc.pb.go")
 		assert.True(t, pathx.FileExists(grpcGo))
@@ -140,12 +140,12 @@ service Greeter {
 			t.Log(err)
 			return
 		}
-		pbDir, err := findPbFile(output, false)
+		pbDir, err := findPbFile(output, protoFile, false)
 		assert.Nil(t, err)
 		pbGo := filepath.Join(pbDir, "greet.pb.go")
 		assert.True(t, pathx.FileExists(pbGo))
 
-		grpcDir, err := findPbFile(output, true)
+		grpcDir, err := findPbFile(output, protoFile, true)
 		assert.Nil(t, err)
 		grpcGo := filepath.Join(grpcDir, "greet_grpc.pb.go")
 		assert.True(t, pathx.FileExists(grpcGo))
@@ -183,12 +183,12 @@ service Greeter {
 			t.Log(err)
 			return
 		}
-		pbDir, err := findPbFile(output, false)
+		pbDir, err := findPbFile(output, protoFile, false)
 		assert.Nil(t, err)
 		pbGo := filepath.Join(pbDir, "greet.pb.go")
 		assert.True(t, pathx.FileExists(pbGo))
 
-		grpcDir, err := findPbFile(output, true)
+		grpcDir, err := findPbFile(output, protoFile, true)
 		assert.Nil(t, err)
 		grpcGo := filepath.Join(grpcDir, "greet_grpc.pb.go")
 		assert.True(t, pathx.FileExists(grpcGo))