Fix StubFunc method to return original stubs (#2)

prashantv · Nov 28, 2016 · 4a17128 · 4a17128
1 parent 6424771
commit 4a17128
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 1 deletion.
diff --git a/func_test.go b/func_test.go
@@ -92,3 +92,27 @@ func TestStubFuncFail(t *testing.T) {
 		}()
 	}
 }
+
+func TestMultipleStubFuncs(t *testing.T) {
+	var f1 = func() int {
+		return 100
+	}
+	var f2 = func() int {
+		return 200
+	}
+	var f3 = func() int {
+		return 300
+	}
+
+	stubs := StubFunc(&f1, 1).StubFunc(&f2, 2)
+	expectVal(t, f1(), 1)
+	expectVal(t, f2(), 2)
+
+	stubs.StubFunc(&f3, 3)
+	expectVal(t, f3(), 3)
+
+	stubs.Reset()
+	expectVal(t, f1(), 100)
+	expectVal(t, f2(), 200)
+	expectVal(t, f3(), 300)
+}
diff --git a/gostub.go b/gostub.go
@@ -68,7 +68,7 @@ func (s *Stubs) StubFunc(funcVarToStub interface{}, stubVal ...interface{}) *Stu
 			funcType.NumOut(), len(stubVal)))
 	}
 
-	return Stub(funcVarToStub, FuncReturning(funcPtrType.Elem(), stubVal...).Interface())
+	return s.Stub(funcVarToStub, FuncReturning(funcPtrType.Elem(), stubVal...).Interface())
 }
 
 // FuncReturning creates a new function with type funcType that returns results.