Merge pull request #72 from JosiahWitt/setup-mocks-ensure-param

Add optional `ensuring.E` parameter to `SetupMocks`

.golangci.yml

@@ -46,3 +46,7 @@ issues:
         - ifshort
         - thelper
         - maintidx
+
+    - text: singleCaseSwitch
+      linters:
+        - gocritic

cmd/ensure/.golangci.yml

@@ -48,6 +48,10 @@ issues:
         - thelper
         - maintidx
 
+    - text: singleCaseSwitch
+      linters:
+        - gocritic
+
     # Exclude some linters from fixtures and scenarios
     - path: .*(/fixtures/)|(/scenarios/).*
       linters:

ensuring/ensuring.go

@@ -14,7 +14,7 @@ import (
 )
 
 //nolint:gochecknoglobals // This is stored as a variable so we can override it for tests in init_test.go.
-var newTestContext = testctx.New
+var newTestContextFunc = testctx.New
 
 // T implements a subset of methods on [testing.T].
 // More methods may be added to T with a minor ensure release.
@@ -118,3 +118,7 @@ func wrap(t T) E {
 		return c
 	}
 }
+
+func newTestContext(t T) testctx.Context {
+	return newTestContextFunc(t, func(t testctx.T) interface{} { return wrap(t) })
+}

ensuring/ensuring_test.go

@@ -188,7 +188,7 @@ func setupMockT(t *testing.T) *mock_testctx.MockT {
 	ctrl := gomock.NewController(t)
 	mockT := mock_testctx.NewMockT(ctrl)
 
-	testhelper.SetTestContext(t, mockT, testctx.New(mockT))
+	testhelper.SetTestContext(t, mockT, testctx.New(mockT, wrapEnsure))
 	return mockT
 }
 
@@ -205,3 +205,7 @@ func setupMockTWithCleanupCheck(t *testing.T) *mock_testctx.MockT {
 
 	return mockT
 }
+
+func wrapEnsure(t testctx.T) interface{} {
+	return ensure.New(t)
+}

ensuring/init_test.go

@@ -5,8 +5,8 @@ import "github.com/JosiahWitt/ensure/ensuring/internal/testhelper"
 
 //nolint:gochecknoinits // Only to make testing easier.
 func init() {
-	// Initializes the unexported newTestContext variable to use the test implementation.
-	// This allows us to continue to keep the tests in the separate testing package and
-	// keep the newTestContext variable unexported.
-	newTestContext = testhelper.NewTestContext
+	// Initializes the unexported newTestContextFunc variable to use the test implementation.
+	// This allows us to continue to keep the tests in the separate testing package and keep
+	// the newTestContextFunc variable unexported.
+	newTestContextFunc = testhelper.NewTestContext
 }

ensuring/internal/testhelper/testhelper.go

@@ -2,6 +2,7 @@
 package testhelper
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/JosiahWitt/ensure/internal/testctx"
@@ -11,15 +12,19 @@ import (
 var (
 	testContexts         = map[testctx.T]testctx.Context{}
 	allowAnyTestContexts = false
+
+	checkingWrapEnsure = false // Used to prevent an infinite loop
 )
 
 // NewTestContext is called instead of [testctx.New] and is setup in ../../init_test.go.
 // This shouldn't be used by anything else.
-func NewTestContext(t testctx.T) testctx.Context {
+func NewTestContext(t testctx.T, wrapEnsure testctx.WrapEnsure) testctx.Context {
+	checkWrapEnsure(t, wrapEnsure)
+
 	ctx, ok := testContexts[t]
 	if !ok {
 		if allowAnyTestContexts {
-			return testctx.New(t)
+			return testctx.New(t, wrapEnsure)
 		}
 
 		panic("Missing mock test context")
@@ -47,3 +52,18 @@ func AllowAnyTestContexts(t *testing.T) {
 		allowAnyTestContexts = false
 	})
 }
+
+func checkWrapEnsure(t testctx.T, wrapEnsure testctx.WrapEnsure) {
+	// Prevent an infinite loop, since NewTestContext will be called by wrapEnsure
+	if checkingWrapEnsure {
+		return
+	}
+
+	checkingWrapEnsure = true
+	defer func() { checkingWrapEnsure = false }()
+
+	ensure := wrapEnsure(t)
+	if ensure == nil || fmt.Sprintf("%T", ensure) != "ensuring.E" {
+		panic(fmt.Sprintf("wrapEnsure doesn't function correctly: %[1]v (%[1]T)", ensure))
+	}
+}

ensuring/run_table_test.go

@@ -69,6 +69,7 @@ func TestERunTableByIndex(t *testing.T) {
 				innerContext := mock_testctx.NewMockContext(ctrl)
 				innerContext.EXPECT().T().Return(innerMockT).AnyTimes()
 				innerContext.EXPECT().GoMockController().Return(gomock.NewController(innerMockT)).AnyTimes()
+				innerContext.EXPECT().Ensure().Return(ensure.New(innerMockT)).AnyTimes()
 				testhelper.SetTestContext(t, innerMockT, innerContext)
 
 				innerMockT.EXPECT().Fatalf(gomock.Any()).Do(func(msg string, args ...interface{}) {
@@ -753,7 +754,7 @@ func (runTableTests) setupMocksField() runTableTestEntryGroup {
 		Prefix: "SetupMocks field",
 		Entries: []runTableTestEntry{
 			{
-				Name:          "with valid function",
+				Name:          "with valid function with one param",
 				ExpectedNames: []string{"name 1", "name 2"},
 				Table: []struct {
 					Name       string
@@ -787,6 +788,44 @@ func (runTableTests) setupMocksField() runTableTestEntryGroup {
 					}
 				},
 			},
+			{
+				Name:                 "with valid function with two params",
+				ExpectedNames:        []string{"name 1", "name 2"},
+				FatalMessagesContain: []string{"first SetupMocks", "second SetupMocks"}, // Not actual failures; only to show ensure is passed in correctly
+				Table: []struct {
+					Name       string
+					Mocks      *TwoValidMocks
+					SetupMocks func(*TwoValidMocks, ensuring.E)
+				}{
+					{
+						Name: "name 1",
+						SetupMocks: func(tvm *TwoValidMocks, ensure ensuring.E) {
+							tvm.Valid1.CustomField = "updated name 1"
+							ensure.Failf("first SetupMocks")
+						},
+					},
+					{
+						Name: "name 2",
+						SetupMocks: func(tvm *TwoValidMocks, ensure ensuring.E) {
+							tvm.Valid1.CustomField = "updated name 2"
+							ensure.Failf("second SetupMocks")
+						},
+					},
+				},
+
+				CheckEntry: func(t *testing.T, rawTable interface{}) {
+					table := rawTable.([]struct {
+						Name       string
+						Mocks      *TwoValidMocks
+						SetupMocks func(*TwoValidMocks, ensuring.E)
+					})
+
+					for _, entry := range table {
+						entry.Mocks.check(t)
+						isTrue(t, entry.Mocks.Valid1.CustomField == "updated "+entry.Name)
+					}
+				},
+			},
 
 			{
 				Name:          "with function not present for one",
@@ -839,7 +878,7 @@ func (runTableTests) setupMocksField() runTableTestEntryGroup {
 
 			{
 				Name:                 "function missing param",
-				FatalMessagesContain: []string{"expected SetupMocks field to be a func(*ensuring_test.TwoValidMocks)"},
+				FatalMessagesContain: []string{"expected SetupMocks field to be one of the following:"},
 				Table: []struct {
 					Name       string
 					Mocks      *TwoValidMocks
@@ -858,7 +897,7 @@ func (runTableTests) setupMocksField() runTableTestEntryGroup {
 
 			{
 				Name:                 "function with invalid param",
-				FatalMessagesContain: []string{"expected SetupMocks field to be a func(*ensuring_test.TwoValidMocks)"},
+				FatalMessagesContain: []string{"expected SetupMocks field to be one of the following:"},
 				Table: []struct {
 					Name       string
 					Mocks      *TwoValidMocks
@@ -877,7 +916,7 @@ func (runTableTests) setupMocksField() runTableTestEntryGroup {
 
 			{
 				Name:                 "function with a return",
-				FatalMessagesContain: []string{"expected SetupMocks field to be a func(*ensuring_test.TwoValidMocks)"},
+				FatalMessagesContain: []string{"expected SetupMocks field to be one of the following:"},
 				Table: []struct {
 					Name       string
 					Mocks      *TwoValidMocks

internal/plugins/internal/id/id.go

@@ -2,6 +2,8 @@
 package id
 
 const (
+	EnsuringE = "ensuring.E"
+
 	Mocks      = "Mocks"
 	SetupMocks = "SetupMocks"
 	Subject    = "Subject"

internal/plugins/setupmocks/setupmocks.go

@@ -2,10 +2,12 @@
 package setupmocks
 
 import (
+	"fmt"
 	"reflect"
 
 	"github.com/JosiahWitt/ensure/internal/plugins"
 	"github.com/JosiahWitt/ensure/internal/plugins/internal/id"
+	"github.com/JosiahWitt/ensure/internal/reflectensure"
 	"github.com/JosiahWitt/ensure/internal/stringerr"
 	"github.com/JosiahWitt/ensure/internal/testctx"
 )
@@ -33,42 +35,67 @@ func (t *TablePlugin) ParseEntryType(entryType reflect.Type) (plugins.TableEntry
 			return nil, stringerr.Newf("%s field must be set on the table to use %s", id.Mocks, id.SetupMocks)
 		}
 
-		if err := validateSetupMocksFieldType(&setupMocksFunc, &mocksStruct); err != nil {
+		funcType, err := parseSetupMocksField(&setupMocksFunc, &mocksStruct)
+		if err != nil {
 			return nil, err
 		}
 
 		h.hasSetupMocks = true
+		h.funcType = funcType
 	}
 
 	return h, nil
 }
 
-func validateSetupMocksFieldType(setupMocksFunc, mocksStruct *reflect.StructField) error {
+func parseSetupMocksField(setupMocksFunc, mocksStruct *reflect.StructField) (funcType, error) {
 	t := setupMocksFunc.Type
 
 	generateError := func() error {
-		return stringerr.Newf("expected %s field to be a func(%v), got: %v", id.SetupMocks, mocksStruct.Type, t)
+		return stringerr.NewBlock(
+			fmt.Sprintf("expected %s field to be one of the following", id.SetupMocks),
+			[]error{
+				stringerr.Newf("func(m %v)", mocksStruct.Type),
+				stringerr.Newf("func(m %v, %s %s)", mocksStruct.Type, id.Ensure, id.EnsuringE),
+			},
+			fmt.Sprintf("Got: %v", t),
+		)
 	}
 
 	if t.Kind() != reflect.Func {
-		return generateError()
+		return 0, generateError()
 	}
 
-	invalidIns := t.NumIn() != 1 || t.In(0) != mocksStruct.Type
-	invalidOuts := t.NumOut() != 0
+	validDefaultIns := t.NumIn() == 1 && t.In(0) == mocksStruct.Type
+	validEnsureIns := t.NumIn() == 2 && t.In(0) == mocksStruct.Type && reflectensure.IsEnsuringE(t.In(1))
+	validIns := validDefaultIns || validEnsureIns
 
-	if invalidIns || invalidOuts {
-		return generateError()
+	validOuts := t.NumOut() == 0
+
+	if !validIns || !validOuts {
+		return 0, generateError()
 	}
 
-	return nil
+	switch {
+	case validEnsureIns:
+		return funcTypeEnsure, nil
+	default:
+		return funcTypeDefault, nil
+	}
 }
 
+type funcType int
+
+const (
+	funcTypeDefault funcType = iota
+	funcTypeEnsure
+)
+
 // TableEntryHooks exposes the before and after hooks for each entry in the table.
 type TableEntryHooks struct {
 	plugins.NoopAfterEntry
 
 	hasSetupMocks bool
+	funcType      funcType
 }
 
 var _ plugins.TableEntryHooks = &TableEntryHooks{}
@@ -88,7 +115,16 @@ func (h *TableEntryHooks) BeforeEntry(ctx testctx.Context, entryValue reflect.Va
 	}
 
 	mocksField := v.FieldByName(id.Mocks)
-	setupMocksFunc.Call([]reflect.Value{mocksField})
+
+	var ins []reflect.Value
+	switch h.funcType {
+	case funcTypeDefault:
+		ins = []reflect.Value{mocksField}
+	case funcTypeEnsure:
+		ins = []reflect.Value{mocksField, reflect.ValueOf(ctx.Ensure())}
+	}
+
+	setupMocksFunc.Call(ins)
 
 	return nil
 }