Allow checking whether all expectations have been met (#21)

Author: FlorianLoch

gomock/callset.go

@@ -18,13 +18,15 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"sync"
 )
 
 // callSet represents a set of expected calls, indexed by receiver and method
 // name.
 type callSet struct {
 	// Calls that are still expected.
-	expected map[callSetKey][]*Call
+	expected   map[callSetKey][]*Call
+	expectedMu *sync.Mutex
 	// Calls that have been exhausted.
 	exhausted map[callSetKey][]*Call
 }
@@ -36,12 +38,20 @@ type callSetKey struct {
 }
 
 func newCallSet() *callSet {
-	return &callSet{make(map[callSetKey][]*Call), make(map[callSetKey][]*Call)}
+	return &callSet{
+		expected:   make(map[callSetKey][]*Call),
+		expectedMu: &sync.Mutex{},
+		exhausted:  make(map[callSetKey][]*Call),
+	}
 }
 
 // Add adds a new expected call.
 func (cs callSet) Add(call *Call) {
 	key := callSetKey{call.receiver, call.method}
+
+	cs.expectedMu.Lock()
+	defer cs.expectedMu.Unlock()
+
 	m := cs.expected
 	if call.exhausted() {
 		m = cs.exhausted
@@ -52,6 +62,10 @@ func (cs callSet) Add(call *Call) {
 // Remove removes an expected call.
 func (cs callSet) Remove(call *Call) {
 	key := callSetKey{call.receiver, call.method}
+
+	cs.expectedMu.Lock()
+	defer cs.expectedMu.Unlock()
+
 	calls := cs.expected[key]
 	for i, c := range calls {
 		if c == call {
@@ -67,6 +81,9 @@ func (cs callSet) Remove(call *Call) {
 func (cs callSet) FindMatch(receiver interface{}, method string, args []interface{}) (*Call, error) {
 	key := callSetKey{receiver, method}
 
+	cs.expectedMu.Lock()
+	defer cs.expectedMu.Unlock()
+
 	// Search through the expected calls.
 	expected := cs.expected[key]
 	var callsErrors bytes.Buffer
@@ -101,6 +118,9 @@ func (cs callSet) FindMatch(receiver interface{}, method string, args []interfac
 
 // Failures returns the calls that are not satisfied.
 func (cs callSet) Failures() []*Call {
+	cs.expectedMu.Lock()
+	defer cs.expectedMu.Unlock()
+
 	failures := make([]*Call, 0, len(cs.expected))
 	for _, calls := range cs.expected {
 		for _, call := range calls {
@@ -111,3 +131,19 @@ func (cs callSet) Failures() []*Call {
 	}
 	return failures
 }
+
+// Satisfied returns true in case all expected calls in this callSet are satisfied.
+func (cs callSet) Satisfied() bool {
+	cs.expectedMu.Lock()
+	defer cs.expectedMu.Unlock()
+
+	for _, calls := range cs.expected {
+		for _, call := range calls {
+			if !call.satisfied() {
+				return false
+			}
+		}
+	}
+
+	return true
+}

gomock/callset_test.go

@@ -77,7 +77,7 @@ func TestCallSetRemove(t *testing.T) {
 
 func TestCallSetFindMatch(t *testing.T) {
 	t.Run("call is exhausted", func(t *testing.T) {
-		cs := callSet{}
+		cs := newCallSet()
 		var receiver interface{} = "TestReceiver"
 		method := "TestMethod"
 		args := []interface{}{}

gomock/controller.go

@@ -231,6 +231,12 @@ func (ctrl *Controller) Finish() {
 	ctrl.finish(false, err)
 }
 
+// Satisfied returns whether all expected calls bound to this Controller have been satisfied.
+// Calling Finish is then guaranteed to not fail due to missing calls.
+func (ctrl *Controller) Satisfied() bool {
+	return ctrl.expectedCalls.Satisfied()
+}
+
 func (ctrl *Controller) finish(cleanup bool, panicErr interface{}) {
 	ctrl.T.Helper()