From 972a8b4b99916619cbc8d49926b1b03485b53935 Mon Sep 17 00:00:00 2001
From: Alancere <804873052@qq.com>
Date: Wed, 13 Dec 2023 15:46:48 +0800
Subject: [PATCH] Added replaceReadmeModule function

---
 .../generator/cmd/v2/common/fileProcessor.go  | 37 +++++++++++++++++++
 .../generator/cmd/v2/common/generation.go     |  5 +++
 2 files changed, 42 insertions(+)

diff --git a/eng/tools/generator/cmd/v2/common/fileProcessor.go b/eng/tools/generator/cmd/v2/common/fileProcessor.go
index 923c9e9fcdcc..1c3241ba67c7 100644
--- a/eng/tools/generator/cmd/v2/common/fileProcessor.go
+++ b/eng/tools/generator/cmd/v2/common/fileProcessor.go
@@ -4,6 +4,7 @@
 package common
 
 import (
+	"bytes"
 	"fmt"
 	"io/fs"
 	"io/ioutil"
@@ -592,3 +593,39 @@ func existSuffixFile(path, suffix string) bool {
 
 	return existed
 }
+
+func replaceReadmeModule(path, rpName, namespaceName, currentVersion string) error {
+	readmeFile, err := os.ReadFile(filepath.Join(path, "README.md"))
+	if err != nil {
+		return err
+	}
+
+	module := fmt.Sprintf("github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/%s/%s", rpName, namespaceName)
+
+	readmeModule := module
+	match := regexp.MustCompile(fmt.Sprintf(`%s/v\d`, module))
+	if match.Match(readmeFile) {
+		readmeModule = match.FindString(string(readmeFile))
+	}
+
+	newModule := module
+	current, err := semver.NewVersion(currentVersion)
+	if err != nil {
+		return err
+	}
+	if current.Major() > 1 {
+		newModule = fmt.Sprintf("%s/v%d", newModule, current.Major())
+	}
+
+	if newModule == readmeModule {
+		return nil
+	}
+
+	newReadmeFile := bytes.ReplaceAll(readmeFile, []byte(readmeModule), []byte(newModule))
+	err = os.WriteFile(filepath.Join(path, "README.md"), newReadmeFile, 0666)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
diff --git a/eng/tools/generator/cmd/v2/common/generation.go b/eng/tools/generator/cmd/v2/common/generation.go
index aa4181a53b41..9f4e716aa3bf 100644
--- a/eng/tools/generator/cmd/v2/common/generation.go
+++ b/eng/tools/generator/cmd/v2/common/generation.go
@@ -328,6 +328,11 @@ func (ctx *GenerateContext) GenerateForSingleRPNamespace(generateParam *Generate
 			}
 		}
 
+		log.Printf("Replace README.md module...")
+		if err = replaceReadmeModule(packagePath, generateParam.RPName, generateParam.NamespaceName, version.String()); err != nil {
+			return nil, err
+		}
+
 		// Example generation should be the last step because the package import relay on the new calculated version
 		if !generateParam.SkipGenerateExample {
 			log.Printf("Generate examples...")