From 70704621de902fa22b84f260d7a7a961efad324d Mon Sep 17 00:00:00 2001
From: RW <rene@gofiber.io>
Date: Sat, 23 Aug 2025 00:55:50 +0200
Subject: [PATCH] fix: add testify import for utils migration

---
 cmd/internal/migrations/v3/common.go      | 23 +++++++++++++++++++++++
 cmd/internal/migrations/v3/common_test.go |  4 ++++
 2 files changed, 27 insertions(+)

diff --git a/cmd/internal/migrations/v3/common.go b/cmd/internal/migrations/v3/common.go
index 4d9668a..751a510 100644
--- a/cmd/internal/migrations/v3/common.go
+++ b/cmd/internal/migrations/v3/common.go
@@ -521,6 +521,26 @@ func MigrateMonitorImport(cmd *cobra.Command, cwd string, _, _ *semver.Version)
 	return nil
 }
 
+func addImport(content, pkg string) string {
+	importStmt := fmt.Sprintf("\"%s\"", pkg)
+	if strings.Contains(content, importStmt) {
+		return content
+	}
+
+	reBlock := regexp.MustCompile(`(?m)^import\s*\(([^)]*)\)`)
+	if reBlock.MatchString(content) {
+		return reBlock.ReplaceAllString(content, fmt.Sprintf("import (\n$1\t%s\n)", importStmt))
+	}
+
+	reSingle := regexp.MustCompile(`(?m)^import\s+([^\n]+)`) // matches single import line
+	if reSingle.MatchString(content) {
+		return reSingle.ReplaceAllString(content, fmt.Sprintf("import (\n\t$1\n\t%s\n)", importStmt))
+	}
+
+	rePkg := regexp.MustCompile(`(?m)^package\s+\w+`)
+	return rePkg.ReplaceAllString(content, fmt.Sprintf("$0\n\nimport (\n\t%s\n)", importStmt))
+}
+
 // MigrateUtilsImport replaces old Fiber utils imports and updates removed helpers
 func MigrateUtilsImport(cmd *cobra.Command, cwd string, _, _ *semver.Version) error {
 	reImport := regexp.MustCompile(`(?m)^(\s*(?:\w+\s+)?)"github\.com/gofiber/fiber/v\d+/utils"$`)
@@ -553,6 +573,9 @@ func MigrateUtilsImport(cmd *cobra.Command, cwd string, _, _ *semver.Version) er
 		content = strings.ReplaceAll(content, "utils.GetBytes", "utils.CopyBytes")
 		content = strings.ReplaceAll(content, "utils.ImmutableString", "string")
 		content = strings.ReplaceAll(content, "utils.AssertEqual", "assert.Equal")
+		if strings.Contains(content, "assert.Equal") {
+			content = addImport(content, "github.com/stretchr/testify/assert")
+		}
 
 		return content
 	})
diff --git a/cmd/internal/migrations/v3/common_test.go b/cmd/internal/migrations/v3/common_test.go
index 350d5e2..5ea3721 100644
--- a/cmd/internal/migrations/v3/common_test.go
+++ b/cmd/internal/migrations/v3/common_test.go
@@ -575,6 +575,7 @@ var (
     _ = utils.GetString([]byte("a"))
     _ = utils.GetBytes("a")
     _ = utils.ImmutableString([]byte("b"))
+    _ = utils.AssertEqual("a", "a")
 )`)
 
 	var buf bytes.Buffer
@@ -583,6 +584,7 @@ var (
 
 	content := readFile(t, file)
 	assert.Contains(t, content, "github.com/gofiber/utils/v2\"")
+	assert.Contains(t, content, "github.com/stretchr/testify/assert")
 	assert.NotContains(t, content, "fiber/v3/utils")
 	assert.NotContains(t, content, "TrimBytes(")
 	assert.NotContains(t, content, "TrimRightBytes(")
@@ -591,6 +593,7 @@ var (
 	assert.NotContains(t, content, "GetString(")
 	assert.NotContains(t, content, "GetBytes(")
 	assert.NotContains(t, content, "ImmutableString(")
+	assert.NotContains(t, content, "utils.AssertEqual")
 	assert.Contains(t, content, "utils.Trim(")
 	assert.Contains(t, content, "utils.TrimRight(")
 	assert.Contains(t, content, "utils.TrimLeft(")
@@ -602,6 +605,7 @@ var (
 	assert.Contains(t, content, "utils.ToString(")
 	assert.Contains(t, content, "utils.CopyBytes(")
 	assert.Contains(t, content, "string([]byte(\"b\"))")
+	assert.Contains(t, content, "assert.Equal")
 	assert.Contains(t, buf.String(), "Migrating utils imports")
 }