Do recover thing for SetArg

Author: msabramo

gomock/call.go

@@ -162,7 +162,8 @@ func (c *Call) Do(f interface{}) *Call {
 				if !ok {
 					panic(r)
 				}
-				if !strings.Contains(errMsg, "reflect: Call using") {
+				if !strings.Contains(errMsg, "reflect: Call using") &&
+					!strings.Contains(errMsg, "reflect.Set: value of") {
 					panic(r)
 				}
 				skipFrames := 8
@@ -277,6 +278,25 @@ func (c *Call) SetArg(n int, value interface{}) *Call {
 		case reflect.Slice:
 			setSlice(args[n], v)
 		default:
+			defer func() {
+				if r := recover(); r != nil {
+					errMsg, ok := r.(string)
+
+					// We only handle a very specific panic
+					// If it's not that one, then we "rethrow" the panic
+					// This allows users to use functions that panic in their tests
+					if !ok {
+						panic(r)
+					}
+					if !strings.Contains(errMsg, "reflect: Call using") &&
+						!strings.Contains(errMsg, "reflect.Set: value of") {
+						panic(r)
+					}
+					skipFrames := 8
+					stackTraceStr := "\n\n" + currentStackTrace(skipFrames)
+					c.t.Fatalf("%s%+v", errMsg, stackTraceStr)
+				}
+			}()
 			reflect.ValueOf(args[n]).Elem().Set(v)
 		}
 		return nil