Add new contextualized API for hooks and steps (#409)

* Add new contextualized API for hooks and steps

* Make default context configurable

* Run AfterStep hooks even after failed steps, fixes #370

* Update CHANGELOG and README

* Add step result status to After hook, fixes #378

* Elaborate hooks documentation

* Add test to pass state between contextualized steps

* Update README with example of passing state between steps

CHANGELOG.md

@@ -12,18 +12,27 @@ This document is formatted according to the principles of [Keep A CHANGELOG](htt
 
 ### Added
 
-* Support for step definitions without return ([364](https://github.com/cucumber/godog/pull/364) -[titouanfreville])
+- Support for step definitions without return ([364](https://github.com/cucumber/godog/pull/364) - [titouanfreville])
+- Contextualized hooks for scenarios and steps ([409](https://github.com/cucumber/godog/pull/409)) - [vearutop])
+- Step result status in After hook ([409](https://github.com/cucumber/godog/pull/409)) - [vearutop])
 
 ### Changed
 
-* Upgraded gherkin-go to v19 ([402](https://github.com/cucumber/godog/pull/402) - [mbow])
+- Upgraded gherkin-go to v19 ([402](https://github.com/cucumber/godog/pull/402) - [mbow])
 
 ### Deprecated
 
+- `ScenarioContext.BeforeScenario`, use `ScenarioContext.Before` ([409](https://github.com/cucumber/godog/pull/409)) - [vearutop])
+- `ScenarioContext.AfterScenario`, use `ScenarioContext.After` ([409](https://github.com/cucumber/godog/pull/409)) - [vearutop])
+- `ScenarioContext.BeforeStep`, use `ScenarioContext.StepContext().Before` ([409](https://github.com/cucumber/godog/pull/409)) - [vearutop])
+- `ScenarioContext.AfterStep`, use `ScenarioContext.StepContext().After` ([409](https://github.com/cucumber/godog/pull/409)) - [vearutop])
+
 ### Removed
 
 ### Fixed
-* Incorrect step definition output for Data Tables ([411](https://github.com/cucumber/godog/pull/411) - [karfrank])
+
+- Incorrect step definition output for Data Tables ([411](https://github.com/cucumber/godog/pull/411) - [karfrank])
+- `ScenarioContext.AfterStep` not invoked after a failed case (([409](https://github.com/cucumber/godog/pull/409)) - [vearutop]))
 
 ## [v0.11.0]
 
@@ -175,7 +184,7 @@ This document is formatted according to the principles of [Keep A CHANGELOG](htt
 
 <!-- Releases -->
 
-[unreleased]: https://github.com/cucumber/godog/compare/v0.11.0...master
+[unreleased]: https://github.com/cucumber/godog/compare/v0.11.0...main
 [v0.11.0]: https://github.com/cucumber/godog/compare/v0.10.0...v0.11.0
 [v0.10.0]: https://github.com/cucumber/godog/compare/v0.9.0...v0.10.0
 [0.9.0]: https://github.com/cucumber/godog/compare/v0.8.1...v0.9.0
@@ -196,3 +205,4 @@ This document is formatted according to the principles of [Keep A CHANGELOG](htt
 [hansbogert]: https://github.com/hansbogert
 [rickardenglund]: https://github.com/rickardenglund
 [mbow]: https://github.com/mbow
+[vearutop]: https://github.com/vearutop

README.md

@@ -248,10 +248,12 @@ godogs
 Now lets implement our step definitions to test our feature requirements:
 
 Replace the contents of `godogs_test.go` with the code below - `vim godogs_test.go`
+
 ```go
 package main
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/cucumber/godog"
@@ -277,25 +279,51 @@ func thereShouldBeRemaining(remaining int) error {
 	return nil
 }
 
-func InitializeTestSuite(ctx *godog.TestSuiteContext) {
+func InitializeTestSuite(sc *godog.TestSuiteContext) {
-	ctx.BeforeSuite(func() { Godogs = 0 })
+	sc.BeforeSuite(func() { Godogs = 0 })
 }
 
-func InitializeScenario(ctx *godog.ScenarioContext) {
+func InitializeScenario(sc *godog.ScenarioContext) {
-	ctx.BeforeScenario(func(*godog.Scenario) {
+	sc.Before(func(ctx context.Context, sc *godog.Scenario) (context.Context, error) {
 		Godogs = 0 // clean the state before every scenario
+
+		return ctx, nil
 	})
 
-	ctx.Step(`^there are (\d+) godogs$`, thereAreGodogs)
+	sc.Step(`^there are (\d+) godogs$`, thereAreGodogs)
-	ctx.Step(`^I eat (\d+)$`, iEat)
+	sc.Step(`^I eat (\d+)$`, iEat)
-	ctx.Step(`^there should be (\d+) remaining$`, thereShouldBeRemaining)
+	sc.Step(`^there should be (\d+) remaining$`, thereShouldBeRemaining)
 }
 ```
 
+You can also pass the state between steps and hooks of a scenario using `context.Context`. 
+Step definitions can receive and return `context.Context`.
+
+```go
+type cntCtxKey struct{} // Key for a particular context value type.
+
+s.Step("^I have a random number of godogs$", func(ctx context.Context) context.Context {
+	// Creating a random number of godog and storing it in context for future reference.
+	cnt := rand.Int()
+	Godogs = cnt
+	return context.WithValue(ctx, cntCtxKey{}, cnt)
+})
+
+s.Step("I eat all available godogs", func(ctx context.Context) error {
+	// Getting previously stored number of godogs from context.
+	cnt := ctx.Value(cntCtxKey{}).(uint32)
+	if Godogs < cnt {
+		return errors.New("can't eat more than I have")
+	}
+	Godogs -= cnt
+	return nil
+})
+```
+
 When you run godog again - `godog`
 
 You should see a passing run:
-```
+```gherkin
 Feature: eat godogs
   In order to be happy
   As a hungry gopher
@@ -305,13 +333,16 @@ Feature: eat godogs
     Given there are 12 godogs        # godogs_test.go:10 -> thereAreGodogs
     When I eat 5                     # godogs_test.go:14 -> iEat
     Then there should be 7 remaining # godogs_test.go:22 -> thereShouldBeRemaining
-
+```
+```
 1 scenarios (1 passed)
 3 steps (3 passed)
 258.302µs
 ```
 
-We have hooked to **BeforeScenario** event in order to reset the application state before each scenario. You may hook into more events, like **AfterStep** to print all state in case of an error. Or **BeforeSuite** to prepare a database.
+We have hooked to `ScenarioContext` **Before** event in order to reset the application state before each scenario. 
+You may hook into more events, like `sc.StepContext()` **After** to print all state in case of an error. 
+Or **BeforeSuite** to prepare a database.
 
 By now, you should have figured out, how to use **godog**. Another advice is to make steps orthogonal, small and simple to read for a user. Whether the user is a dumb website user or an API developer, who may understand a little more technical context - it should target that user.
 
@@ -369,7 +400,7 @@ var opts = godog.Options{
 
 func init() {
 	godog.BindFlags("godog.", pflag.CommandLine, &opts) // godog v0.10.0 and earlier
-	godog.BindCommandLineFlags("godog.", &opts)        // godog v0.11.0 (latest)
+	godog.BindCommandLineFlags("godog.", &opts)        // godog v0.11.0 and later
 }
 
 func TestMain(m *testing.M) {

features/events.feature

@@ -72,6 +72,9 @@ Feature: suite events
 
         Scenario: two
           Then passing step
+          And failing step
+          And adding step state to context
+          And having correct context
 
         Scenario Outline: three
           Then passing step
@@ -84,7 +87,9 @@ Feature: suite events
     Then these events had to be fired for a number of times:
       | BeforeSuite    | 1 |
       | BeforeScenario | 2 |
-      | BeforeStep     | 2 |
+      | BeforeStep     | 5 |
-      | AfterStep      | 2 |
+      | AfterStep      | 5 |
       | AfterScenario  | 2 |
       | AfterSuite     | 1 |
+
+    And the suite should have failed

internal/flags/options.go

@@ -1,6 +1,7 @@
 package flags
 
 import (
+	"context"
 	"io"
 )
 
@@ -56,4 +57,7 @@ type Options struct {
 
 	// Where it should print formatter output
 	Output io.Writer
+
+	// DefaultContext is used as initial context instead of context.Background().
+	DefaultContext context.Context
 }

internal/models/stepdef.go

@@ -1,6 +1,7 @@
 package models
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"reflect"
@@ -32,91 +33,104 @@ type StepDefinition struct {
 	Undefined []string
 }
 
+var typeOfContext = reflect.TypeOf((*context.Context)(nil)).Elem()
+
 // Run a step with the matched arguments using reflect
-func (sd *StepDefinition) Run() interface{} {
+func (sd *StepDefinition) Run(ctx context.Context) (context.Context, interface{}) {
+	var values []reflect.Value
+
 	typ := sd.HandlerValue.Type()
-	if len(sd.Args) < typ.NumIn() {
+	numIn := typ.NumIn()
-		return fmt.Errorf("%w: expected %d arguments, matched %d from step", ErrUnmatchedStepArgumentNumber, typ.NumIn(), len(sd.Args))
+	hasCtxIn := numIn > 0 && typ.In(0).Implements(typeOfContext)
+	ctxOffset := 0
+
+	if hasCtxIn {
+		values = append(values, reflect.ValueOf(ctx))
+		ctxOffset = 1
+		numIn--
 	}
 
-	var values []reflect.Value
+	if len(sd.Args) < numIn {
-	for i := 0; i < typ.NumIn(); i++ {
+		return ctx, fmt.Errorf("%w: expected %d arguments, matched %d from step", ErrUnmatchedStepArgumentNumber, typ.NumIn(), len(sd.Args))
-		param := typ.In(i)
+	}
+
+	for i := 0; i < numIn; i++ {
+		param := typ.In(i + ctxOffset)
 		switch param.Kind() {
 		case reflect.Int:
 			s, err := sd.shouldBeString(i)
 			if err != nil {
-				return err
+				return ctx, err
 			}
 			v, err := strconv.ParseInt(s, 10, 0)
 			if err != nil {
-				return fmt.Errorf(`%w %d: "%s" to int: %s`, ErrCannotConvert, i, s, err)
+				return ctx, fmt.Errorf(`%w %d: "%s" to int: %s`, ErrCannotConvert, i, s, err)
 			}
 			values = append(values, reflect.ValueOf(int(v)))
 		case reflect.Int64:
 			s, err := sd.shouldBeString(i)
 			if err != nil {
-				return err
+				return ctx, err
 			}
 			v, err := strconv.ParseInt(s, 10, 64)
 			if err != nil {
-				return fmt.Errorf(`%w %d: "%s" to int64: %s`, ErrCannotConvert, i, s, err)
+				return ctx, fmt.Errorf(`%w %d: "%s" to int64: %s`, ErrCannotConvert, i, s, err)
 			}
 			values = append(values, reflect.ValueOf(v))
 		case reflect.Int32:
 			s, err := sd.shouldBeString(i)
 			if err != nil {
-				return err
+				return ctx, err
 			}
 			v, err := strconv.ParseInt(s, 10, 32)
 			if err != nil {
-				return fmt.Errorf(`%w %d: "%s" to int32: %s`, ErrCannotConvert, i, s, err)
+				return ctx, fmt.Errorf(`%w %d: "%s" to int32: %s`, ErrCannotConvert, i, s, err)
 			}
 			values = append(values, reflect.ValueOf(int32(v)))
 		case reflect.Int16:
 			s, err := sd.shouldBeString(i)
 			if err != nil {
-				return err
+				return ctx, err
 			}
 			v, err := strconv.ParseInt(s, 10, 16)
 			if err != nil {
-				return fmt.Errorf(`%w %d: "%s" to int16: %s`, ErrCannotConvert, i, s, err)
+				return ctx, fmt.Errorf(`%w %d: "%s" to int16: %s`, ErrCannotConvert, i, s, err)
 			}
 			values = append(values, reflect.ValueOf(int16(v)))
 		case reflect.Int8:
 			s, err := sd.shouldBeString(i)
 			if err != nil {
-				return err
+				return ctx, err
 			}
 			v, err := strconv.ParseInt(s, 10, 8)
 			if err != nil {
-				return fmt.Errorf(`%w %d: "%s" to int8: %s`, ErrCannotConvert, i, s, err)
+				return ctx, fmt.Errorf(`%w %d: "%s" to int8: %s`, ErrCannotConvert, i, s, err)
 			}
 			values = append(values, reflect.ValueOf(int8(v)))
 		case reflect.String:
 			s, err := sd.shouldBeString(i)
 			if err != nil {
-				return err
+				return ctx, err
 			}
 			values = append(values, reflect.ValueOf(s))
 		case reflect.Float64:
 			s, err := sd.shouldBeString(i)
 			if err != nil {
-				return err
+				return ctx, err
 			}
 			v, err := strconv.ParseFloat(s, 64)
 			if err != nil {
-				return fmt.Errorf(`%w %d: "%s" to float64: %s`, ErrCannotConvert, i, s, err)
+				return ctx, fmt.Errorf(`%w %d: "%s" to float64: %s`, ErrCannotConvert, i, s, err)
 			}
 			values = append(values, reflect.ValueOf(v))
 		case reflect.Float32:
 			s, err := sd.shouldBeString(i)
 			if err != nil {
-				return err
+				return ctx, err
 			}
 			v, err := strconv.ParseFloat(s, 32)
 			if err != nil {
-				return fmt.Errorf(`%w %d: "%s" to float32: %s`, ErrCannotConvert, i, s, err)
+				return ctx, fmt.Errorf(`%w %d: "%s" to float32: %s`, ErrCannotConvert, i, s, err)
 			}
 			values = append(values, reflect.ValueOf(float32(v)))
 		case reflect.Ptr:
@@ -133,7 +147,7 @@ func (sd *StepDefinition) Run() interface{} {
 					break
 				}
 
-				return fmt.Errorf(`%w %d: "%v" of type "%T" to *messages.PickleDocString`, ErrCannotConvert, i, arg, arg)
+				return ctx, fmt.Errorf(`%w %d: "%v" of type "%T" to *messages.PickleDocString`, ErrCannotConvert, i, arg, arg)
 			case "messages.PickleTable":
 				if v, ok := arg.(*messages.PickleStepArgument); ok {
 					values = append(values, reflect.ValueOf(v.DataTable))
@@ -145,32 +159,42 @@ func (sd *StepDefinition) Run() interface{} {
 					break
 				}
 
-				return fmt.Errorf(`%w %d: "%v" of type "%T" to *messages.PickleTable`, ErrCannotConvert, i, arg, arg)
+				return ctx, fmt.Errorf(`%w %d: "%v" of type "%T" to *messages.PickleTable`, ErrCannotConvert, i, arg, arg)
 			default:
-				return fmt.Errorf("%w: the argument %d type %T is not supported %s", ErrUnsupportedArgumentType, i, arg, param.Elem().String())
+				return ctx, fmt.Errorf("%w: the argument %d type %T is not supported %s", ErrUnsupportedArgumentType, i, arg, param.Elem().String())
 			}
 		case reflect.Slice:
 			switch param {
 			case typeOfBytes:
 				s, err := sd.shouldBeString(i)
 				if err != nil {
-					return err
+					return ctx, err
 				}
 				values = append(values, reflect.ValueOf([]byte(s)))
 			default:
-				return fmt.Errorf("%w: the slice argument %d type %s is not supported", ErrUnsupportedArgumentType, i, param.Kind())
+				return ctx, fmt.Errorf("%w: the slice argument %d type %s is not supported", ErrUnsupportedArgumentType, i, param.Kind())
 			}
 		default:
-			return fmt.Errorf("%w: the argument %d type %s is not supported", ErrUnsupportedArgumentType, i, param.Kind())
+			return ctx, fmt.Errorf("%w: the argument %d type %s is not supported", ErrUnsupportedArgumentType, i, param.Kind())
 		}
 	}
 
 	res := sd.HandlerValue.Call(values)
 	if len(res) == 0 {
-		return nil
+		return ctx, nil
 	}
 
-	return res[0].Interface()
+	if len(res) == 1 {
+		r := res[0].Interface()
+
+		if ctx, ok := r.(context.Context); ok {
+			return ctx, nil
+		}
+
+		return ctx, res[0].Interface()
+	}
+
+	return res[0].Interface().(context.Context), res[1].Interface()
 }
 
 func (sd *StepDefinition) shouldBeString(idx int) (string, error) {

internal/models/stepdef_test.go

@@ -1,6 +1,7 @@
 package models_test
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"reflect"
@@ -8,12 +9,59 @@ import (
 	"testing"
 
 	"github.com/cucumber/messages-go/v16"
+	"github.com/stretchr/testify/assert"
 
 	"github.com/cucumber/godog"
 	"github.com/cucumber/godog/formatters"
 	"github.com/cucumber/godog/internal/models"
 )
 
+func TestShouldSupportContext(t *testing.T) {
+	ctx := context.WithValue(context.Background(), "original", 123)
+
+	fn := func(ctx context.Context, a int64, b int32, c int16, d int8) context.Context {
+		assert.Equal(t, 123, ctx.Value("original"))
+
+		return context.WithValue(ctx, "updated", 321)
+	}
+
+	def := &models.StepDefinition{
+		StepDefinition: formatters.StepDefinition{
+			Handler: fn,
+		},
+		HandlerValue: reflect.ValueOf(fn),
+	}
+
+	def.Args = []interface{}{"1", "1", "1", "1"}
+	ctx, err := def.Run(ctx)
+	assert.Nil(t, err)
+	assert.Equal(t, 123, ctx.Value("original"))
+	assert.Equal(t, 321, ctx.Value("updated"))
+}
+
+func TestShouldSupportContextAndError(t *testing.T) {
+	ctx := context.WithValue(context.Background(), "original", 123)
+
+	fn := func(ctx context.Context, a int64, b int32, c int16, d int8) (context.Context, error) {
+		assert.Equal(t, 123, ctx.Value("original"))
+
+		return context.WithValue(ctx, "updated", 321), nil
+	}
+
+	def := &models.StepDefinition{
+		StepDefinition: formatters.StepDefinition{
+			Handler: fn,
+		},
+		HandlerValue: reflect.ValueOf(fn),
+	}
+
+	def.Args = []interface{}{"1", "1", "1", "1"}
+	ctx, err := def.Run(ctx)
+	assert.Nil(t, err)
+	assert.Equal(t, 123, ctx.Value("original"))
+	assert.Equal(t, 321, ctx.Value("updated"))
+}
+
 func TestShouldSupportEmptyHandlerReturn(t *testing.T) {
 	fn := func(a int64, b int32, c int16, d int8) {}
 
@@ -25,12 +73,12 @@ func TestShouldSupportEmptyHandlerReturn(t *testing.T) {
 	}
 
 	def.Args = []interface{}{"1", "1", "1", "1"}
-	if err := def.Run(); err != nil {
+	if _, err := def.Run(context.Background()); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
 	def.Args = []interface{}{"1", "1", "1", strings.Repeat("1", 9)}
-	if err := def.Run(); err == nil {
+	if _, err := def.Run(context.Background()); err == nil {
 		t.Fatalf("expected convertion fail for int8, but got none")
 	}
 }
@@ -46,12 +94,12 @@ func TestShouldSupportIntTypes(t *testing.T) {
 	}
 
 	def.Args = []interface{}{"1", "1", "1", "1"}
-	if err := def.Run(); err != nil {
+	if _, err := def.Run(context.Background()); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
 	def.Args = []interface{}{"1", "1", "1", strings.Repeat("1", 9)}
-	if err := def.Run(); err == nil {
+	if _, err := def.Run(context.Background()); err == nil {
 		t.Fatalf("expected convertion fail for int8, but got none")
 	}
 }
@@ -67,12 +115,12 @@ func TestShouldSupportFloatTypes(t *testing.T) {
 	}
 
 	def.Args = []interface{}{"1.1", "1.09"}
-	if err := def.Run(); err != nil {
+	if _, err := def.Run(context.Background()); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
 	def.Args = []interface{}{"1.08", strings.Repeat("1", 65) + ".67"}
-	if err := def.Run(); err == nil {
+	if _, err := def.Run(context.Background()); err == nil {
 		t.Fatalf("expected convertion fail for float32, but got none")
 	}
 }
@@ -104,15 +152,15 @@ func TestShouldNotSupportOtherPointerTypesThanGherkin(t *testing.T) {
 		Args:         []interface{}{(*messages.PickleTable)(nil)},
 	}
 
-	if err := def1.Run(); err == nil {
+	if _, err := def1.Run(context.Background()); err == nil {
 		t.Fatalf("expected conversion error, but got none")
 	}
 
-	if err := def2.Run(); err != nil {
+	if _, err := def2.Run(context.Background()); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
-	if err := def3.Run(); err != nil {
+	if _, err := def3.Run(context.Background()); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 }
@@ -136,11 +184,11 @@ func TestShouldSupportOnlyByteSlice(t *testing.T) {
 		Args:         []interface{}{[]string{}},
 	}
 
-	if err := def1.Run(); err != nil {
+	if _, err := def1.Run(context.Background()); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
-	if err := def2.Run(); err == nil {
+	if _, err := def2.Run(context.Background()); err == nil {
 		t.Fatalf("expected conversion error, but got none")
 	}
 }
@@ -156,7 +204,7 @@ func TestUnexpectedArguments(t *testing.T) {
 
 	def.Args = []interface{}{"1"}
 
-	res := def.Run()
+	_, res := def.Run(context.Background())
 	if res == nil {
 		t.Fatalf("expected an error due to wrong number of arguments, but got none")
 	}
@@ -182,7 +230,7 @@ func TestStepDefinition_Run_StepShouldBeString(t *testing.T) {
 
 		def.Args = []interface{}{12}
 
-		res := def.Run()
+		_, res := def.Run(context.Background())
 		if res == nil {
 			t.Fatalf("expected a string convertion error, but got none")
 		}
@@ -224,7 +272,7 @@ func TestStepDefinition_Run_InvalidHandlerParamConversion(t *testing.T) {
 
 		def.Args = []interface{}{12}
 
-		res := def.Run()
+		_, res := def.Run(context.Background())
 		if res == nil {
 			t.Fatalf("expected an unsupported argument type error, but got none")
 		}
@@ -280,7 +328,7 @@ func TestStepDefinition_Run_StringConversionToFunctionType(t *testing.T) {
 			Args:         args,
 		}
 
-		res := def.Run()
+		_, res := def.Run(context.Background())
 		if res == nil {
 			t.Fatalf("expected a cannot convert argument type error, but got none")
 		}
@@ -321,7 +369,7 @@ func TestStepDefinition_Run_StringConversionToFunctionType(t *testing.T) {
 // def = &models.StepDefinition{Handler: fn2, HandlerValue: reflect.ValueOf(fn2)}
 
 // def.Args = []interface{}{"1"}
-// if err := def.Run(); err == nil {
+// if _, err := def.Run(context.Background()); err == nil {
 // 	t.Fatalf("expected an error due to wrong argument type, but got none")
 // }
 
@@ -347,7 +395,7 @@ func TestShouldSupportDocStringToStringConversion(t *testing.T) {
 		}},
 	}
 
-	if err := def.Run(); err != nil {
+	if _, err := def.Run(context.Background()); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 }

run.go

@@ -1,6 +1,7 @@
 package godog
 
 import (
+	"context"
 	"fmt"
 	"go/build"
 	"io"
@@ -36,6 +37,8 @@ type runner struct {
 	randomSeed            int64
 	stopOnFailure, strict bool
 
+	defaultContext context.Context
+
 	features []*models.Feature
 
 	testSuiteInitializer testSuiteInitializer
@@ -102,6 +105,7 @@ func (r *runner) concurrent(rate int) (failed bool) {
 					randomSeed:     r.randomSeed,
 					strict:         r.strict,
 					storage:        r.storage,
+					defaultContext: r.defaultContext,
 				}
 
 				if r.scenarioInitializer != nil {
@@ -231,6 +235,7 @@ func runWithOptions(suiteName string, runner runner, opt Options) int {
 
 	runner.stopOnFailure = opt.StopOnFailure
 	runner.strict = opt.Strict
+	runner.defaultContext = opt.DefaultContext
 
 	// store chosen seed in environment, so it could be seen in formatter summary report
 	os.Setenv("GODOG_SEED", strconv.FormatInt(runner.randomSeed, 10))

run_test.go

@@ -414,11 +414,11 @@ func Test_AllFeaturesRun(t *testing.T) {
 ...................................................................... 140
 ...................................................................... 210
 ...................................................................... 280
-............................                                           308
+.............................                                          309
 
 
 81 scenarios (81 passed)
-308 steps (308 passed)
+309 steps (309 passed)
 0s
 `
 

suite.go

@@ -1,6 +1,7 @@
 package godog
 
 import (
+	"context"
 	"fmt"
 	"reflect"
 	"strings"
@@ -13,7 +14,10 @@ import (
 	"github.com/cucumber/godog/internal/utils"
 )
 
-var errorInterface = reflect.TypeOf((*error)(nil)).Elem()
+var (
+	errorInterface   = reflect.TypeOf((*error)(nil)).Elem()
+	contextInterface = reflect.TypeOf((*context.Context)(nil)).Elem()
+)
 
 // ErrUndefined is returned in case if step definition was not found
 var ErrUndefined = fmt.Errorf("step is undefined")
@@ -22,6 +26,22 @@ var ErrUndefined = fmt.Errorf("step is undefined")
 // step implementation is pending
 var ErrPending = fmt.Errorf("step implementation is pending")
 
+// StepResultStatus describes step result.
+type StepResultStatus = models.StepResultStatus
+
+const (
+	// StepPassed indicates step that passed.
+	StepPassed StepResultStatus = models.Passed
+	// StepFailed indicates step that failed.
+	StepFailed = models.Failed
+	// StepSkipped indicates step that was skipped.
+	StepSkipped = models.Skipped
+	// StepUndefined indicates undefined step.
+	StepUndefined = models.Undefined
+	// StepPending indicates step with pending implementation.
+	StepPending = models.Pending
+)
+
 type suite struct {
 	steps []*models.StepDefinition
 
@@ -33,11 +53,13 @@ type suite struct {
 	stopOnFailure bool
 	strict        bool
 
+	defaultContext context.Context
+
 	// suite event handlers
-	beforeScenarioHandlers []func(*Scenario)
+	beforeScenarioHandlers []BeforeScenarioHook
-	beforeStepHandlers     []func(*Step)
+	beforeStepHandlers     []BeforeStepHook
-	afterStepHandlers      []func(*Step, error)
+	afterStepHandlers      []AfterStepHook
-	afterScenarioHandlers  []func(*Scenario, error)
+	afterScenarioHandlers  []AfterScenarioHook
 }
 
 func (s *suite) matchStep(step *messages.PickleStep) *models.StepDefinition {
@@ -48,15 +70,11 @@ func (s *suite) matchStep(step *messages.PickleStep) *models.StepDefinition {
 	return def
 }
 
-func (s *suite) runStep(pickle *messages.Pickle, step *messages.PickleStep, prevStepErr error) (err error) {
+func (s *suite) runStep(ctx context.Context, pickle *Scenario, step *Step, prevStepErr error) (rctx context.Context, err error) {
-	// run before step handlers
+	var (
-	for _, f := range s.beforeStepHandlers {
+		match *models.StepDefinition
-		f(step)
+		sr    = models.PickleStepResult{Status: models.Undefined}
-	}
+	)
-
-	match := s.matchStep(step)
-	s.storage.MustInsertStepDefintionMatch(step.AstNodeIds[0], match)
-	s.fmt.Defined(pickle, step, match.GetInternalStepDefinition())
 
 	// user multistep definitions may panic
 	defer func() {
@@ -67,6 +85,24 @@ func (s *suite) runStep(pickle *messages.Pickle, step *messages.PickleStep, prev
 			}
 		}
 
+		defer func() {
+			// run after step handlers
+			for _, f := range s.afterStepHandlers {
+				hctx, herr := f(rctx, step, sr.Status, err)
+
+				// Adding hook error to resulting error without breaking hooks loop.
+				if herr != nil {
+					if err == nil {
+						err = herr
+					} else {
+						err = fmt.Errorf("%v: %w", herr, err)
+					}
+				}
+
+				rctx = hctx
+			}
+		}()
+
 		if prevStepErr != nil {
 			return
 		}
@@ -75,7 +111,7 @@ func (s *suite) runStep(pickle *messages.Pickle, step *messages.PickleStep, prev
 			return
 		}
 
-		sr := models.NewStepResult(pickle.Id, step.Id, match)
+		sr = models.NewStepResult(pickle.Id, step.Id, match)
 
 		switch err {
 		case nil:
@@ -95,15 +131,22 @@ func (s *suite) runStep(pickle *messages.Pickle, step *messages.PickleStep, prev
 
 			s.fmt.Failed(pickle, step, match.GetInternalStepDefinition(), err)
 		}
-
-		// run after step handlers
-		for _, f := range s.afterStepHandlers {
-			f(step, err)
-		}
 	}()
 
-	if undef, err := s.maybeUndefined(step.Text, step.Argument); err != nil {
+	// run before step handlers
-		return err
+	for _, f := range s.beforeStepHandlers {
+		ctx, err = f(ctx, step)
+		if err != nil {
+			return ctx, err
+		}
+	}
+
+	match = s.matchStep(step)
+	s.storage.MustInsertStepDefintionMatch(step.AstNodeIds[0], match)
+	s.fmt.Defined(pickle, step, match.GetInternalStepDefinition())
+
+	if ctx, undef, err := s.maybeUndefined(ctx, step.Text, step.Argument); err != nil {
+		return ctx, err
 	} else if len(undef) > 0 {
 		if match != nil {
 			match = &models.StepDefinition{
@@ -118,82 +161,85 @@ func (s *suite) runStep(pickle *messages.Pickle, step *messages.PickleStep, prev
 			}
 		}
 
-		sr := models.NewStepResult(pickle.Id, step.Id, match)
+		sr = models.NewStepResult(pickle.Id, step.Id, match)
 		sr.Status = models.Undefined
 		s.storage.MustInsertPickleStepResult(sr)
 
 		s.fmt.Undefined(pickle, step, match.GetInternalStepDefinition())
-		return ErrUndefined
+		return ctx, ErrUndefined
 	}
 
 	if prevStepErr != nil {
-		sr := models.NewStepResult(pickle.Id, step.Id, match)
+		sr = models.NewStepResult(pickle.Id, step.Id, match)
 		sr.Status = models.Skipped
 		s.storage.MustInsertPickleStepResult(sr)
 
 		s.fmt.Skipped(pickle, step, match.GetInternalStepDefinition())
-		return nil
+		return ctx, nil
 	}
 
-	err = s.maybeSubSteps(match.Run())
+	ctx, err = s.maybeSubSteps(match.Run(ctx))
-	return
+
+	return ctx, err
 }
 
-func (s *suite) maybeUndefined(text string, arg interface{}) ([]string, error) {
+func (s *suite) maybeUndefined(ctx context.Context, text string, arg interface{}) (context.Context, []string, error) {
 	step := s.matchStepText(text)
 	if nil == step {
-		return []string{text}, nil
+		return ctx, []string{text}, nil
 	}
 
 	var undefined []string
 	if !step.Nested {
-		return undefined, nil
+		return ctx, undefined, nil
 	}
 
 	if arg != nil {
 		step.Args = append(step.Args, arg)
 	}
 
-	for _, next := range step.Run().(Steps) {
+	ctx, steps := step.Run(ctx)
+
+	for _, next := range steps.(Steps) {
 		lines := strings.Split(next, "\n")
 		// @TODO: we cannot currently parse table or content body from nested steps
 		if len(lines) > 1 {
-			return undefined, fmt.Errorf("nested steps cannot be multiline and have table or content body argument")
+			return ctx, undefined, fmt.Errorf("nested steps cannot be multiline and have table or content body argument")
 		}
 		if len(lines[0]) > 0 && lines[0][len(lines[0])-1] == ':' {
-			return undefined, fmt.Errorf("nested steps cannot be multiline and have table or content body argument")
+			return ctx, undefined, fmt.Errorf("nested steps cannot be multiline and have table or content body argument")
 		}
-		undef, err := s.maybeUndefined(next, nil)
+		ctx, undef, err := s.maybeUndefined(ctx, next, nil)
 		if err != nil {
-			return undefined, err
+			return ctx, undefined, err
 		}
 		undefined = append(undefined, undef...)
 	}
-	return undefined, nil
+	return ctx, undefined, nil
 }
 
-func (s *suite) maybeSubSteps(result interface{}) error {
+func (s *suite) maybeSubSteps(ctx context.Context, result interface{}) (context.Context, error) {
 	if nil == result {
-		return nil
+		return ctx, nil
 	}
 
 	if err, ok := result.(error); ok {
-		return err
+		return ctx, err
 	}
 
 	steps, ok := result.(Steps)
 	if !ok {
-		return fmt.Errorf("unexpected error, should have been []string: %T - %+v", result, result)
+		return ctx, fmt.Errorf("unexpected error, should have been []string: %T - %+v", result, result)
 	}
 
 	for _, text := range steps {
 		if def := s.matchStepText(text); def == nil {
-			return ErrUndefined
+			return ctx, ErrUndefined
-		} else if err := s.maybeSubSteps(def.Run()); err != nil {
+		} else if ctx, err := s.maybeSubSteps(def.Run(ctx)); err != nil {
-			return fmt.Errorf("%s: %+v", text, err)
+			return ctx, fmt.Errorf("%s: %+v", text, err)
 		}
 	}
-	return nil
+	return ctx, nil
 }
 
 func (s *suite) matchStepText(text string) *models.StepDefinition {
@@ -220,9 +266,13 @@ func (s *suite) matchStepText(text string) *models.StepDefinition {
 	return nil
 }
 
-func (s *suite) runSteps(pickle *messages.Pickle, steps []*messages.PickleStep) (err error) {
+func (s *suite) runSteps(ctx context.Context, pickle *Scenario, steps []*Step) (context.Context, error) {
+	var (
+		stepErr, err error
+	)
+
 	for _, step := range steps {
-		stepErr := s.runStep(pickle, step, err)
+		ctx, stepErr = s.runStep(ctx, pickle, step, err)
 		switch stepErr {
 		case ErrUndefined:
 			// do not overwrite failed error
@@ -236,7 +286,8 @@ func (s *suite) runSteps(pickle *messages.Pickle, steps []*messages.PickleStep)
 			err = stepErr
 		}
 	}
-	return
+
+	return ctx, err
 }
 
 func (s *suite) shouldFail(err error) bool {
@@ -262,6 +313,11 @@ func isEmptyFeature(pickles []*messages.Pickle) bool {
 }
 
 func (s *suite) runPickle(pickle *messages.Pickle) (err error) {
+	ctx := s.defaultContext
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
 	if len(pickle.Steps) == 0 {
 		pr := models.PickleResult{PickleID: pickle.Id, StartedAt: utils.TimeNowFunc()}
 		s.storage.MustInsertPickleResult(pr)
@@ -272,7 +328,10 @@ func (s *suite) runPickle(pickle *messages.Pickle) (err error) {
 
 	// run before scenario handlers
 	for _, f := range s.beforeScenarioHandlers {
-		f(pickle)
+		ctx, err = f(ctx, pickle)
+		if err != nil {
+			return err
+		}
 	}
 
 	pr := models.PickleResult{PickleID: pickle.Id, StartedAt: utils.TimeNowFunc()}
@@ -281,12 +340,23 @@ func (s *suite) runPickle(pickle *messages.Pickle) (err error) {
 	s.fmt.Pickle(pickle)
 
 	// scenario
-	err = s.runSteps(pickle, pickle.Steps)
+	ctx, err = s.runSteps(ctx, pickle, pickle.Steps)
 
 	// run after scenario handlers
 	for _, f := range s.afterScenarioHandlers {
-		f(pickle, err)
+		hctx, herr := f(ctx, pickle, err)
+
+		// Adding hook error to resulting error without breaking hooks loop.
+		if herr != nil {
+			if err == nil {
+				err = herr
+			} else {
+				err = fmt.Errorf("%v: %w", herr, err)
+			}
 		}
 
-	return
+		ctx = hctx
+	}
+
+	return err
 }

suite_context_test.go

@@ -2,8 +2,10 @@ package godog
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"encoding/xml"
+	"errors"
 	"fmt"
 	"path/filepath"
 	"regexp"
@@ -34,7 +36,7 @@ import (
 func InitializeScenario(ctx *ScenarioContext) {
 	tc := &godogFeaturesScenario{}
 
-	ctx.BeforeScenario(tc.ResetBeforeEachScenario)
+	ctx.Before(tc.ResetBeforeEachScenario)
 
 	ctx.Step(`^(?:a )?feature path "([^"]*)"$`, tc.featurePath)
 	ctx.Step(`^I parse features$`, tc.parseFeatures)
@@ -108,19 +110,42 @@ func InitializeScenario(ctx *ScenarioContext) {
 		return nil
 	})
 
-	ctx.Step(`^(?:a )?passing step without return$`, func() {})
+	ctx.Step(`^passing step without return$`, func() {})
-	ctx.BeforeStep(tc.inject)
+
+	ctx.Step(`^having correct context$`, func(ctx context.Context) (context.Context, error) {
+		if ctx.Value(ctxKey("BeforeScenario")) == nil {
+			return ctx, errors.New("missing BeforeScenario in context")
 		}
 
-func (tc *godogFeaturesScenario) inject(step *Step) {
+		if ctx.Value(ctxKey("BeforeStep")) == nil {
+			return ctx, errors.New("missing BeforeStep in context")
+		}
+
+		if ctx.Value(ctxKey("StepState")) == nil {
+			return ctx, errors.New("missing StepState in context")
+		}
+
+		return context.WithValue(ctx, ctxKey("Step"), true), nil
+	})
+
+	ctx.Step(`^adding step state to context$`, func(ctx context.Context) context.Context {
+		return context.WithValue(ctx, ctxKey("StepState"), true)
+	})
+
+	ctx.StepContext().Before(tc.inject)
+}
+
+type ctxKey string
+
+func (tc *godogFeaturesScenario) inject(ctx context.Context, step *Step) (context.Context, error) {
 	if !tc.allowInjection {
-		return
+		return ctx, nil
 	}
 
 	step.Text = injectAll(step.Text)
 
 	if step.Argument == nil {
-		return
+		return ctx, nil
 	}
 
 	if table := step.Argument.DataTable; table != nil {
@@ -134,6 +159,8 @@ func (tc *godogFeaturesScenario) inject(step *Step) {
 	if doc := step.Argument.DocString; doc != nil {
 		doc.Content = injectAll(doc.Content)
 	}
+
+	return ctx, nil
 }
 
 func injectAll(src string) string {
@@ -167,7 +194,7 @@ type godogFeaturesScenario struct {
 	allowInjection   bool
 }
 
-func (tc *godogFeaturesScenario) ResetBeforeEachScenario(*Scenario) {
+func (tc *godogFeaturesScenario) ResetBeforeEachScenario(ctx context.Context, sc *Scenario) (context.Context, error) {
 	// reset whole suite with the state
 	tc.out.Reset()
 	tc.paths = []string{}
@@ -179,6 +206,8 @@ func (tc *godogFeaturesScenario) ResetBeforeEachScenario(*Scenario) {
 	// reset all fired events
 	tc.events = []*firedEvent{}
 	tc.allowInjection = false
+
+	return ctx, nil
 }
 
 func (tc *godogFeaturesScenario) iSetVariableInjectionTo(to string) error {
@@ -391,20 +420,56 @@ func (tc *godogFeaturesScenario) iAmListeningToSuiteEvents() error {
 
 	scenarioContext := ScenarioContext{suite: tc.testedSuite}
 
-	scenarioContext.BeforeScenario(func(pickle *Scenario) {
+	scenarioContext.Before(func(ctx context.Context, pickle *Scenario) (context.Context, error) {
 		tc.events = append(tc.events, &firedEvent{"BeforeScenario", []interface{}{pickle}})
+
+		return context.WithValue(ctx, ctxKey("BeforeScenario"), true), nil
 	})
 
-	scenarioContext.AfterScenario(func(pickle *Scenario, err error) {
+	scenarioContext.After(func(ctx context.Context, pickle *Scenario, err error) (context.Context, error) {
 		tc.events = append(tc.events, &firedEvent{"AfterScenario", []interface{}{pickle, err}})
+
+		if ctx.Value(ctxKey("BeforeScenario")) == nil {
+			return ctx, errors.New("missing BeforeScenario in context")
+		}
+
+		if ctx.Value(ctxKey("AfterStep")) == nil {
+			return ctx, errors.New("missing AfterStep in context")
+		}
+
+		return context.WithValue(ctx, ctxKey("AfterScenario"), true), nil
 	})
 
-	scenarioContext.BeforeStep(func(step *Step) {
+	scenarioContext.StepContext().Before(func(ctx context.Context, step *Step) (context.Context, error) {
 		tc.events = append(tc.events, &firedEvent{"BeforeStep", []interface{}{step}})
+
+		if ctx.Value(ctxKey("BeforeScenario")) == nil {
+			return ctx, errors.New("missing BeforeScenario in context")
+		}
+
+		return context.WithValue(ctx, ctxKey("BeforeStep"), true), nil
 	})
 
-	scenarioContext.AfterStep(func(step *Step, err error) {
+	scenarioContext.StepContext().After(func(ctx context.Context, step *Step, status StepResultStatus, err error) (context.Context, error) {
 		tc.events = append(tc.events, &firedEvent{"AfterStep", []interface{}{step, err}})
+
+		if ctx.Value(ctxKey("BeforeScenario")) == nil {
+			return ctx, errors.New("missing BeforeScenario in context")
+		}
+
+		if ctx.Value(ctxKey("BeforeStep")) == nil {
+			return ctx, errors.New("missing BeforeStep in context")
+		}
+
+		if step.Text == "having correct context" && ctx.Value(ctxKey("Step")) == nil {
+			if status != StepSkipped {
+				return ctx, fmt.Errorf("unexpected step result status: %s", status)
+			}
+
+			return ctx, errors.New("missing Step in context")
+		}
+
+		return context.WithValue(ctx, ctxKey("AfterStep"), true), nil
 	})
 
 	return nil

test_context.go

@@ -1,6 +1,7 @@
 package godog
 
 import (
+	"context"
 	"fmt"
 	"reflect"
 	"regexp"
@@ -12,9 +13,6 @@ import (
 	"github.com/cucumber/godog/internal/models"
 )
 
-// matchable errors
-var ()
-
 // Scenario represents the executed scenario
 type Scenario = messages.Pickle
 
@@ -97,26 +95,101 @@ type ScenarioContext struct {
 	suite *suite
 }
 
+// StepContext allows registering step hooks.
+type StepContext struct {
+	suite *suite
+}
+
+// Before registers a a function or method
+// to be run before every scenario.
+//
+// It is a good practice to restore the default state
+// before every scenario so it would be isolated from
+// any kind of state.
+func (ctx ScenarioContext) Before(h BeforeScenarioHook) {
+	ctx.suite.beforeScenarioHandlers = append(ctx.suite.beforeScenarioHandlers, h)
+}
+
+// BeforeScenarioHook defines a hook before scenario.
+type BeforeScenarioHook func(ctx context.Context, sc *Scenario) (context.Context, error)
+
+// After registers an function or method
+// to be run after every scenario.
+func (ctx ScenarioContext) After(h AfterScenarioHook) {
+	ctx.suite.afterScenarioHandlers = append(ctx.suite.afterScenarioHandlers, h)
+}
+
+// AfterScenarioHook defines a hook after scenario.
+type AfterScenarioHook func(ctx context.Context, sc *Scenario, err error) (context.Context, error)
+
+// StepContext exposes StepContext of a scenario.
+func (ctx *ScenarioContext) StepContext() StepContext {
+	return StepContext{suite: ctx.suite}
+}
+
+// Before registers a function or method
+// to be run before every step.
+func (ctx StepContext) Before(h BeforeStepHook) {
+	ctx.suite.beforeStepHandlers = append(ctx.suite.beforeStepHandlers, h)
+}
+
+// BeforeStepHook defines a hook before step.
+type BeforeStepHook func(ctx context.Context, st *Step) (context.Context, error)
+
+// After registers an function or method
+// to be run after every step.
+//
+// It may be convenient to return a different kind of error
+// in order to print more state details which may help
+// in case of step failure
+//
+// In some cases, for example when running a headless
+// browser, to take a screenshot after failure.
+func (ctx StepContext) After(h AfterStepHook) {
+	ctx.suite.afterStepHandlers = append(ctx.suite.afterStepHandlers, h)
+}
+
+// AfterStepHook defines a hook after step.
+type AfterStepHook func(ctx context.Context, st *Step, status StepResultStatus, err error) (context.Context, error)
+
 // BeforeScenario registers a function or method
 // to be run before every scenario.
 //
 // It is a good practice to restore the default state
 // before every scenario so it would be isolated from
 // any kind of state.
+//
+// Deprecated: use Before.
 func (ctx *ScenarioContext) BeforeScenario(fn func(sc *Scenario)) {
-	ctx.suite.beforeScenarioHandlers = append(ctx.suite.beforeScenarioHandlers, fn)
+	ctx.Before(func(ctx context.Context, sc *Scenario) (context.Context, error) {
+		fn(sc)
+
+		return ctx, nil
+	})
 }
 
 // AfterScenario registers an function or method
 // to be run after every scenario.
+//
+// Deprecated: use After.
 func (ctx *ScenarioContext) AfterScenario(fn func(sc *Scenario, err error)) {
-	ctx.suite.afterScenarioHandlers = append(ctx.suite.afterScenarioHandlers, fn)
+	ctx.After(func(ctx context.Context, sc *Scenario, err error) (context.Context, error) {
+		fn(sc, err)
+
+		return ctx, nil
+	})
 }
 
 // BeforeStep registers a function or method
 // to be run before every step.
+//
+// Deprecated: use ScenarioContext.StepContext() and StepContext.Before.
 func (ctx *ScenarioContext) BeforeStep(fn func(st *Step)) {
-	ctx.suite.beforeStepHandlers = append(ctx.suite.beforeStepHandlers, fn)
+	ctx.StepContext().Before(func(ctx context.Context, st *Step) (context.Context, error) {
+		fn(st)
+
+		return ctx, nil
+	})
 }
 
 // AfterStep registers an function or method
@@ -128,8 +201,14 @@ func (ctx *ScenarioContext) BeforeStep(fn func(st *Step)) {
 //
 // In some cases, for example when running a headless
 // browser, to take a screenshot after failure.
+//
+// Deprecated: use ScenarioContext.StepContext() and StepContext.After.
 func (ctx *ScenarioContext) AfterStep(fn func(st *Step, err error)) {
-	ctx.suite.afterStepHandlers = append(ctx.suite.afterStepHandlers, fn)
+	ctx.StepContext().After(func(ctx context.Context, st *Step, status StepResultStatus, err error) (context.Context, error) {
+		fn(st, err)
+
+		return ctx, nil
+	})
 }
 
 // Step allows to register a *StepDefinition in the
@@ -179,8 +258,8 @@ func (ctx *ScenarioContext) Step(expr, stepFunc interface{}) {
 		panic(fmt.Sprintf("expected handler to be func, but got: %T", stepFunc))
 	}
 
-	if typ.NumOut() > 1 {
+	if typ.NumOut() > 2 {
-		panic(fmt.Sprintf("expected handler to return either zero or one value, but it has: %d", typ.NumOut()))
+		panic(fmt.Sprintf("expected handler to return either zero, one or two values, but it has: %d", typ.NumOut()))
 	}
 
 	def := &models.StepDefinition{
@@ -195,8 +274,8 @@ func (ctx *ScenarioContext) Step(expr, stepFunc interface{}) {
 		typ = typ.Out(0)
 		switch typ.Kind() {
 		case reflect.Interface:
-			if !typ.Implements(errorInterface) {
+			if !typ.Implements(errorInterface) && !typ.Implements(contextInterface) {
-				panic(fmt.Sprintf("expected handler to return an error, but got: %s", typ.Kind()))
+				panic(fmt.Sprintf("expected handler to return an error or context.Context, but got: %s", typ.Kind()))
 			}
 		case reflect.Slice:
 			if typ.Elem().Kind() != reflect.String {

test_context_test.go

@@ -41,13 +41,13 @@ func TestScenarioContext_Step(t *testing.T) {
 				So(func() { ctx.Step(".*", 124) }, ShouldPanicWith, fmt.Sprintf("expected handler to be func, but got: %T", 12))
 			})
 
-			Convey("has more than 1 return value", func() {
+			Convey("has more than 2 return values", func() {
-				So(func() { ctx.Step(".*", nokLimitCase) }, ShouldPanicWith, fmt.Sprintf("expected handler to return either zero or one value, but it has: 2"))
+				So(func() { ctx.Step(".*", nokLimitCase) }, ShouldPanicWith, fmt.Sprintf("expected handler to return either zero, one or two values, but it has: 3"))
-				So(func() { ctx.Step(".*", nokMore) }, ShouldPanicWith, fmt.Sprintf("expected handler to return either zero or one value, but it has: 5"))
+				So(func() { ctx.Step(".*", nokMore) }, ShouldPanicWith, fmt.Sprintf("expected handler to return either zero, one or two values, but it has: 5"))
 			})
 
 			Convey("return type is not an error or string slice or void", func() {
-				So(func() { ctx.Step(".*", nokInvalidReturnInterfaceType) }, ShouldPanicWith, "expected handler to return an error, but got: interface")
+				So(func() { ctx.Step(".*", nokInvalidReturnInterfaceType) }, ShouldPanicWith, "expected handler to return an error or context.Context, but got: interface")
 				So(func() { ctx.Step(".*", nokInvalidReturnSliceType) }, ShouldPanicWith, "expected handler to return []string for multistep, but got: []int")
 				So(func() { ctx.Step(".*", nokInvalidReturnOtherType) }, ShouldPanicWith, "expected handler to return an error or []string, but got: chan")
 			})
@@ -60,7 +60,7 @@ func TestScenarioContext_Step(t *testing.T) {
 func okEmptyResult()                             {}
 func okErrorResult() error                       { return nil }
 func okSliceResult() []string                    { return nil }
-func nokLimitCase() (int, error)                 { return 0, nil }
+func nokLimitCase() (string, int, error)         { return "", 0, nil }
 func nokMore() (int, int, int, int, error)       { return 0, 0, 0, 0, nil }
 func nokInvalidReturnInterfaceType() interface{} { return 0 }
 func nokInvalidReturnSliceType() []int           { return nil }