fix import error if generate multiple proto (zeromicro#3694)

rosas99 · Nov 4, 2023 · df2799f · df2799f
1 parent fd8ee0b
commit df2799f
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 15 deletions.
diff --git 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
@@ -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))