From c1076ab37f88e47cf087274f68ae3de9a5be4a6d Mon Sep 17 00:00:00 2001
From: Jordan Krage <jmank88@gmail.com>
Date: Wed, 18 Dec 2024 12:25:32 -0600
Subject: [PATCH] mock: update callString to format context.Context as pointer

---
 mock/mock.go      |  6 ++++++
 mock/mock_test.go | 19 +++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/mock/mock.go b/mock/mock.go
index eb5682df9..99ef59c52 100644
--- a/mock/mock.go
+++ b/mock/mock.go
@@ -1,6 +1,7 @@
 package mock
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"path"
@@ -446,6 +447,11 @@ func callString(method string, arguments Arguments, includeArgumentValues bool)
 	if includeArgumentValues {
 		var argVals []string
 		for argIndex, arg := range arguments {
+			if ctx, ok := arg.(context.Context); ok {
+				// avoid data race from formating context directly
+				argVals = append(argVals, fmt.Sprintf("%d: %p", argIndex, ctx))
+				continue
+			}
 			if _, ok := arg.(*FunctionalOptionsArgument); ok {
 				argVals = append(argVals, fmt.Sprintf("%d: %s", argIndex, arg))
 				continue
diff --git a/mock/mock_test.go b/mock/mock_test.go
index 04664d44b..32de4b081 100644
--- a/mock/mock_test.go
+++ b/mock/mock_test.go
@@ -1,6 +1,7 @@
 package mock
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"regexp"
@@ -1234,6 +1235,24 @@ func Test_callString(t *testing.T) {
 	assert.Equal(t, `Method(int,bool,string)`, callString("Method", []interface{}{1, true, "something"}, false))
 	assert.Equal(t, `Method(<nil>)`, callString("Method", []interface{}{nil}, false))
 
+	t.Run("context", func(t *testing.T) {
+		// special treatment for context.Context to avoid data races
+		ctx, cancel := context.WithCancel(context.Background())
+		t.Cleanup(cancel)
+		done := make(chan struct{})
+		t.Cleanup(func() {
+			<-done
+		})
+		go func() {
+			defer close(done)
+			for i := 0; i < 20; i++ {
+				cancel()
+			}
+		}()
+		assert.Equal(t, fmt.Sprintf(`Method(*context.cancelCtx)
+		0: %p`, ctx),
+			callString("Method", []interface{}{ctx}, true))
+	})
 }
 
 func Test_Mock_Called(t *testing.T) {