From bcf5ec509f891bcb9bbe85287b0053cd274b8670 Mon Sep 17 00:00:00 2001
From: RW <rene@gofiber.io>
Date: Fri, 1 Aug 2025 15:31:58 +0200
Subject: [PATCH] Fix context migration and tests

---
 cmd/internal/migrations/v3/common.go      | 23 +++++++++++------
 cmd/internal/migrations/v3/common_test.go | 30 +++++++++++++++++++++--
 2 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/cmd/internal/migrations/v3/common.go b/cmd/internal/migrations/v3/common.go
index eff1dd7..3520b9b 100644
--- a/cmd/internal/migrations/v3/common.go
+++ b/cmd/internal/migrations/v3/common.go
@@ -98,14 +98,23 @@ func MigrateGenericHelpers(cmd *cobra.Command, cwd string, _, _ *semver.Version)
 
 // MigrateContextMethods updates context related methods to the new names
 func MigrateContextMethods(cmd *cobra.Command, cwd string, _, _ *semver.Version) error {
-	replacer := strings.NewReplacer(
-		".Context()", ".RequestCtx()",
-		".UserContext()", ".Context()",
-		".SetUserContext(", ".SetContext(", // TODO: check if this is correct
-	)
-
 	err := internal.ChangeFileContent(cwd, func(content string) string {
-		return replacer.Replace(content)
+		// UserContext() removed - Ctx implements context.Context
+		reUserCtx := regexp.MustCompile(`(\w+)\.UserContext\(\)`)
+		content = reUserCtx.ReplaceAllString(content, `$1`)
+
+		// SetUserContext removed - comment out the call
+		reSetUserCtx := regexp.MustCompile(`(?m)^(\s*)(.*\.SetUserContext\([^\n]*\).*)$`)
+		content = reSetUserCtx.ReplaceAllString(content, `$1// TODO: SetUserContext was removed, please migrate manually: $2`)
+
+		// old Context() returned fasthttp.RequestCtx
+		reReqCtx := regexp.MustCompile(`(\w+)\.Context\(\)`)
+		content = reReqCtx.ReplaceAllString(content, `$1.RequestCtx()`)
+
+		// remaining Context() usages return context.Context -> remove call
+		content = strings.ReplaceAll(content, ".Context()", "")
+
+		return content
 	})
 	if err != nil {
 		return fmt.Errorf("failed to migrate context methods: %w", err)
diff --git a/cmd/internal/migrations/v3/common_test.go b/cmd/internal/migrations/v3/common_test.go
index a97d404..0ad06c4 100644
--- a/cmd/internal/migrations/v3/common_test.go
+++ b/cmd/internal/migrations/v3/common_test.go
@@ -171,8 +171,34 @@ func handler(c fiber.Ctx) error {
 
 	content := readFile(t, file)
 	assert.Contains(t, content, ".RequestCtx()")
-	assert.Contains(t, content, ".Context()")
-	assert.Contains(t, content, ".SetContext(")
+	assert.NotContains(t, content, ".Context()")
+	assert.Contains(t, content, `// TODO: SetUserContext was removed, please migrate manually: c.SetUserContext(ctx)`)
+	assert.Contains(t, content, "uc := c")
+	assert.Contains(t, buf.String(), "Migrating context methods")
+}
+
+func Test_MigrateContextMethods_SetUserContextAssignment(t *testing.T) {
+	t.Parallel()
+
+	dir, err := os.MkdirTemp("", "mcmtest2")
+	require.NoError(t, err)
+	defer func() { require.NoError(t, os.RemoveAll(dir)) }()
+
+	file := writeTempFile(t, dir, `package main
+import "github.com/gofiber/fiber/v2"
+func handler(ctx fiber.Ctx) error {
+    res := ctx.SetUserContext(ctx.Context())
+    _ = res
+    return nil
+}`)
+
+	var buf bytes.Buffer
+	cmd := newCmd(&buf)
+	require.NoError(t, v3.MigrateContextMethods(cmd, dir, nil, nil))
+
+	content := readFile(t, file)
+	assert.Contains(t, content, `// TODO: SetUserContext was removed, please migrate manually: res := ctx.SetUserContext(ctx.RequestCtx())`)
+	assert.NotContains(t, content, `.UserContext()`)
 	assert.Contains(t, buf.String(), "Migrating context methods")
 }