allow screenshots to be taken via test code for all hooks, #632

execution/merge.go

@@ -39,6 +39,8 @@ func mergeDataTableSpecResults(sResult *result.SuiteResult) *result.SuiteResult
 	suiteRes.Tags = sResult.Tags
 	suiteRes.PreHookMessages = append(suiteRes.PreHookMessages, sResult.PreHookMessages...)
 	suiteRes.PostHookMessages = append(suiteRes.PostHookMessages, sResult.PostHookMessages...)
+	suiteRes.PreHookScreenshots = append(suiteRes.PreHookScreenshots, sResult.PreHookScreenshots...)
+	suiteRes.PostHookScreenshots = append(suiteRes.PostHookScreenshots, sResult.PostHookScreenshots...)
 	combinedResults := make(map[string][]*result.SpecResult)
 	for _, res := range sResult.SpecResults {
 		fileName := res.ProtoSpec.GetFileName()

execution/result/suiteResult.go

@@ -42,6 +42,8 @@ type SuiteResult struct {
 	SpecsSkippedCount int
 	PreHookMessages   []string
 	PostHookMessages  []string
+	PreHookScreenshots   [][]byte
+	PostHookScreenshots  [][]byte
 }
 
 // NewSuiteResult is a constructor for SuitResult

execution/scenarioExecutor.go

@@ -122,6 +122,7 @@ func (e *scenarioExecutor) notifyBeforeScenarioHook(scenarioResult *result.Scena
 	e.pluginHandler.NotifyPlugins(message)
 	res := executeHook(message, scenarioResult, e.runner)
 	scenarioResult.ProtoScenario.PreHookMessages = res.Message
+	scenarioResult.ProtoScenario.PreHookScreenshots = res.ScreenShot
 	if res.GetFailed() {
 		setScenarioFailure(e.currentExecutionInfo)
 		handleHookFailure(scenarioResult, res, result.AddPreHook)
@@ -133,6 +134,7 @@ func (e *scenarioExecutor) notifyAfterScenarioHook(scenarioResult *result.Scenar
 		ScenarioExecutionEndingRequest: &gauge_messages.ScenarioExecutionEndingRequest{CurrentExecutionInfo: e.currentExecutionInfo}}
 	res := executeHook(message, scenarioResult, e.runner)
 	scenarioResult.ProtoScenario.PostHookMessages = res.Message
+	scenarioResult.ProtoScenario.PostHookScreenshots = res.ScreenShot
 	if res.GetFailed() {
 		setScenarioFailure(e.currentExecutionInfo)
 		handleHookFailure(scenarioResult, res, result.AddPostHook)

execution/scenarioExecutor_test.go

@@ -87,3 +87,73 @@ func TestNotifyAfterScenarioShouldAddAfterScenarioHookMessages(t *testing.T) {
 		t.Errorf("Expected `After Scenario Called` message, got : %s", gotMessages[0])
 	}
 }
+
+func TestNotifyBeforeScenarioShouldAddBeforeScenarioHookScreenshots(t *testing.T) {
+	r := &mockRunner{}
+	h := &mockPluginHandler{NotifyPluginsfunc: func(m *gauge_messages.Message) {}, GracefullyKillPluginsfunc: func() {}}
+	r.ExecuteAndGetStatusFunc = func(m *gauge_messages.Message) *gauge_messages.ProtoExecutionResult {
+		if m.MessageType == gauge_messages.Message_ScenarioExecutionStarting {
+			return &gauge_messages.ProtoExecutionResult{
+				ScreenShot:    [][]byte{[]byte("screenshot1"), []byte("screenshot2")},
+				Failed:        false,
+				ExecutionTime: 10,
+			}
+		}
+		return &gauge_messages.ProtoExecutionResult{}
+	}
+	ei := &gauge_messages.ExecutionInfo{}
+	sce := newScenarioExecutor(r, h, ei, nil, nil, nil, 0)
+	scenario := &gauge.Scenario{
+		Heading: &gauge.Heading{Value: "A scenario"},
+		Span:    &gauge.Span{Start: 2, End: 10},
+	}
+	scenarioResult := result.NewScenarioResult(gauge.NewProtoScenario(scenario))
+	sce.notifyBeforeScenarioHook(scenarioResult)
+	beforeScenarioScreenShots := scenarioResult.ProtoScenario.PreHookScreenshots
+	expected := []string{"screenshot1", "screenshot2"}
+
+	if len(beforeScenarioScreenShots) != len(expected) {
+		t.Errorf("Expected 2 screenshots, got : %d", len(beforeScenarioScreenShots))
+	}
+
+	for i, e := range expected {
+		if string(beforeScenarioScreenShots[i]) != e {
+			t.Errorf("Expected `%s` screenshot, got : %s", e, beforeScenarioScreenShots[i])
+		}
+	}
+}
+
+func TestNotifyAfterScenarioShouldAddAfterScenarioHookScreenshots(t *testing.T) {
+	r := &mockRunner{}
+	h := &mockPluginHandler{NotifyPluginsfunc: func(m *gauge_messages.Message) {}, GracefullyKillPluginsfunc: func() {}}
+	r.ExecuteAndGetStatusFunc = func(m *gauge_messages.Message) *gauge_messages.ProtoExecutionResult {
+		if m.MessageType == gauge_messages.Message_ScenarioExecutionEnding {
+			return &gauge_messages.ProtoExecutionResult{
+				ScreenShot:    [][]byte{[]byte("screenshot1"), []byte("screenshot2")},
+				Failed:        false,
+				ExecutionTime: 10,
+			}
+		}
+		return &gauge_messages.ProtoExecutionResult{}
+	}
+	ei := &gauge_messages.ExecutionInfo{}
+	sce := newScenarioExecutor(r, h, ei, nil, nil, nil, 0)
+	scenario := &gauge.Scenario{
+		Heading: &gauge.Heading{Value: "A scenario"},
+		Span:    &gauge.Span{Start: 2, End: 10},
+	}
+	scenarioResult := result.NewScenarioResult(gauge.NewProtoScenario(scenario))
+	sce.notifyAfterScenarioHook(scenarioResult)
+	afterScenarioScreenShots := scenarioResult.ProtoScenario.PostHookScreenshots
+	expected := []string{"screenshot1", "screenshot2"}
+
+	if len(afterScenarioScreenShots) != len(expected) {
+		t.Errorf("Expected 2 screenshots, got : %d", len(afterScenarioScreenShots))
+	}
+
+	for i, e := range expected {
+		if string(afterScenarioScreenShots[i]) != e {
+			t.Errorf("Expected `%s` screenshot, got : %s", e, afterScenarioScreenShots[i])
+		}
+	}
+}

execution/simpleExecution.go

@@ -150,6 +150,7 @@ func (e *simpleExecution) notifyBeforeSuite() {
 		ExecutionStartingRequest: &gauge_messages.ExecutionStartingRequest{}}
 	res := e.executeHook(m)
 	e.suiteResult.PreHookMessages = res.Message
+	e.suiteResult.PreHookScreenshots = res.ScreenShot
 	if res.GetFailed() {
 		handleHookFailure(e.suiteResult, res, result.AddPreHook)
 	}
@@ -160,6 +161,7 @@ func (e *simpleExecution) notifyAfterSuite() {
 		ExecutionEndingRequest: &gauge_messages.ExecutionEndingRequest{CurrentExecutionInfo: e.currentExecutionInfo}}
 	res := e.executeHook(m)
 	e.suiteResult.PostHookMessages = res.Message
+	e.suiteResult.PostHookScreenshots = res.ScreenShot
 	if res.GetFailed() {
 		handleHookFailure(e.suiteResult, res, result.AddPostHook)
 	}

execution/simpleExecution_test.go

@@ -80,3 +80,65 @@ func TestNotifyAfterSuiteShouldAddsAfterSuiteHookMessages(t *testing.T) {
 		t.Errorf("Expected `After Suite Called` message, got : %s", gotMessages[0])
 	}
 }
+
+func TestNotifyBeforeSuiteShouldAddsBeforeSuiteHookScreenshots(t *testing.T) {
+	r := &mockRunner{}
+	h := &mockPluginHandler{NotifyPluginsfunc: func(m *gauge_messages.Message) {}, GracefullyKillPluginsfunc: func() {}}
+	r.ExecuteAndGetStatusFunc = func(m *gauge_messages.Message) *gauge_messages.ProtoExecutionResult {
+		if m.MessageType == gauge_messages.Message_ExecutionStarting {
+			return &gauge_messages.ProtoExecutionResult{
+				ScreenShot:    [][]byte{[]byte("screenshot1"), []byte("screenshot2")},
+				Failed:        false,
+				ExecutionTime: 10,
+			}
+		}
+		return &gauge_messages.ProtoExecutionResult{}
+	}
+	ei := &executionInfo{runner: r, pluginHandler: h}
+	simpleExecution := newSimpleExecution(ei, false)
+	simpleExecution.suiteResult = result.NewSuiteResult(ExecuteTags, simpleExecution.startTime)
+	simpleExecution.notifyBeforeSuite()
+
+	beforeSuiteScreenshots := simpleExecution.suiteResult.PreHookScreenshots
+	expected := []string{"screenshot1", "screenshot2"}
+
+	if len(beforeSuiteScreenshots) != len(expected) {
+		t.Errorf("Expected 2 screenshots, got : %d", len(beforeSuiteScreenshots))
+	}
+	for i, e := range expected {
+		if string(beforeSuiteScreenshots[i]) != e {
+			t.Errorf("Expected `%s` screenshot, got : %s", e, beforeSuiteScreenshots[i])
+		}
+	}
+}
+
+func TestNotifyAfterSuiteShouldAddsAfterSuiteHookScreenshots(t *testing.T) {
+	r := &mockRunner{}
+	h := &mockPluginHandler{NotifyPluginsfunc: func(m *gauge_messages.Message) {}, GracefullyKillPluginsfunc: func() {}}
+	r.ExecuteAndGetStatusFunc = func(m *gauge_messages.Message) *gauge_messages.ProtoExecutionResult {
+		if m.MessageType == gauge_messages.Message_ExecutionEnding {
+			return &gauge_messages.ProtoExecutionResult{
+				ScreenShot:    [][]byte{[]byte("screenshot1"), []byte("screenshot2")},
+				Failed:        false,
+				ExecutionTime: 10,
+			}
+		}
+		return &gauge_messages.ProtoExecutionResult{}
+	}
+	ei := &executionInfo{runner: r, pluginHandler: h}
+	simpleExecution := newSimpleExecution(ei, false)
+	simpleExecution.suiteResult = result.NewSuiteResult(ExecuteTags, simpleExecution.startTime)
+	simpleExecution.notifyAfterSuite()
+
+	afterSuiteScreenshots := simpleExecution.suiteResult.PostHookScreenshots
+	expected := []string{"screenshot1", "screenshot2"}
+
+	if len(afterSuiteScreenshots) != len(expected) {
+		t.Errorf("Expected 2 screenshots, got : %d", len(afterSuiteScreenshots))
+	}
+	for i, e := range expected {
+		if string(afterSuiteScreenshots[i]) != e {
+			t.Errorf("Expected `%s` screenshot, got : %s", e, afterSuiteScreenshots[i])
+		}
+	}
+}

execution/specExecutor.go

@@ -244,6 +244,7 @@ func (e *specExecutor) notifyBeforeSpecHook() {
 	e.pluginHandler.NotifyPlugins(m)
 	res := executeHook(m, e.specResult, e.runner)
 	e.specResult.ProtoSpec.PreHookMessages = res.Message
+	e.specResult.ProtoSpec.PreHookScreenshots = res.ScreenShot
 	if res.GetFailed() {
 		setSpecFailure(e.currentExecutionInfo)
 		handleHookFailure(e.specResult, res, result.AddPreHook)
@@ -256,6 +257,7 @@ func (e *specExecutor) notifyAfterSpecHook() {
 		SpecExecutionEndingRequest: &gauge_messages.SpecExecutionEndingRequest{CurrentExecutionInfo: e.currentExecutionInfo}}
 	res := executeHook(m, e.specResult, e.runner)
 	e.specResult.ProtoSpec.PostHookMessages = res.Message
+	e.specResult.ProtoSpec.PostHookScreenshots = res.ScreenShot
 	if res.GetFailed() {
 		setSpecFailure(e.currentExecutionInfo)
 		handleHookFailure(e.specResult, res, result.AddPostHook)

execution/specExecutor_test.go

@@ -688,6 +688,53 @@ func TestExecuteAddsSpecHookExecutionMessages(t *testing.T) {
 	}
 }
 
+func TestExecuteAddsSpecHookExecutionScreenshots(t *testing.T) {
+	errs := gauge.NewBuildErrors()
+	mockRunner := &mockRunner{}
+	mockHandler := &mockPluginHandler{NotifyPluginsfunc: func(m *gauge_messages.Message) {}, GracefullyKillPluginsfunc: func() {}}
+
+	mockRunner.ExecuteAndGetStatusFunc = func(m *gauge_messages.Message) *gauge_messages.ProtoExecutionResult {
+		if m.MessageType == gauge_messages.Message_SpecExecutionEnding {
+			return &gauge_messages.ProtoExecutionResult{
+				ScreenShot:    [][]byte{[]byte("screenshot1"), []byte("screenshot2")},
+				Failed:        false,
+				ExecutionTime: 10,
+			}
+		} else if m.MessageType == gauge_messages.Message_SpecExecutionStarting {
+			return &gauge_messages.ProtoExecutionResult{
+				ScreenShot:    [][]byte{[]byte("screenshot3"), []byte("screenshot4")},
+				Failed:        false,
+				ExecutionTime: 10,
+			}
+		}
+		return &gauge_messages.ProtoExecutionResult{}
+	}
+	se := newSpecExecutor(exampleSpec, mockRunner, mockHandler, errs, 0)
+	se.execute(true, false, true)
+
+	beforeSpecScreenshots := se.specResult.ProtoSpec.PreHookScreenshots
+	afterSpecScreenshots := se.specResult.ProtoSpec.PostHookScreenshots
+	expectedAfterSpecScreenshots := []string{"screenshot1", "screenshot2"}
+	expectedBeforeSpecScreenshots := []string{"screenshot3", "screenshot4"}
+
+	if len(beforeSpecScreenshots) != len(expectedBeforeSpecScreenshots) {
+		t.Errorf("Expected 2 screenshots, got : %d", len(beforeSpecScreenshots))
+	}
+	for i, e := range expectedBeforeSpecScreenshots {
+		if string(beforeSpecScreenshots[i]) != e {
+			t.Errorf("Expected `%s` screenshot, got : %s", e, beforeSpecScreenshots[i])
+		}
+	}
+	if len(afterSpecScreenshots) != len(expectedAfterSpecScreenshots) {
+		t.Errorf("Expected 2 screenshots, got : %d", len(afterSpecScreenshots))
+	}
+	for i, e := range expectedAfterSpecScreenshots {
+		if string(afterSpecScreenshots[i]) != e {
+			t.Errorf("Expected `%s` screenshot, got : %s", e, afterSpecScreenshots[i])
+		}
+	}
+}
+
 func TestExecuteShouldNotifyAfterSpecEvent(t *testing.T) {
 	errs := gauge.NewBuildErrors()
 	r := &mockRunner{}

execution/stepExecutor.go

@@ -76,6 +76,7 @@ func (e *stepExecutor) notifyBeforeStepHook(stepResult *result.StepResult) {
 	e.pluginHandler.NotifyPlugins(m)
 	res := executeHook(m, stepResult, e.runner)
 	stepResult.ProtoStep.PreHookMessages = res.Message
+	stepResult.ProtoStep.PreHookScreenshots = res.ScreenShot
 	if res.GetFailed() {
 		setStepFailure(e.currentExecutionInfo)
 		handleHookFailure(stepResult, res, result.AddPreHook)
@@ -90,7 +91,7 @@ func (e *stepExecutor) notifyAfterStepHook(stepResult *result.StepResult) {
 
 	res := executeHook(m, stepResult, e.runner)
 	stepResult.ProtoStep.PostHookMessages = append(stepResult.ProtoStepExecResult().GetExecutionResult().Message, res.Message...)
-	stepResult.ProtoStepExecResult().GetExecutionResult().ScreenShot = append(stepResult.ProtoStepExecResult().GetExecutionResult().ScreenShot, res.ScreenShot...)
+	stepResult.ProtoStep.PostHookScreenshots = append(stepResult.ProtoStepExecResult().GetExecutionResult().ScreenShot, res.ScreenShot...)
 	if res.GetFailed() {
 		setStepFailure(e.currentExecutionInfo)
 		handleHookFailure(stepResult, res, result.AddPostHook)

execution/stepExecutor_test.go

@@ -111,7 +111,55 @@ func TestStepExecutionShouldAddAfterStepHookMessages(t *testing.T) {
 	}
 }
 
-func TestStepExecutionShouldGetScreenshots(t *testing.T) {
+func TestStepExecutionShouldGetScreenshotsBeforeStep(t *testing.T) {
+	r := &mockRunner{}
+	h := &mockPluginHandler{NotifyPluginsfunc: func(m *gauge_messages.Message) {}, GracefullyKillPluginsfunc: func() {}}
+	r.ExecuteAndGetStatusFunc = func(m *gauge_messages.Message) *gauge_messages.ProtoExecutionResult {
+		if m.MessageType == gauge_messages.Message_StepExecutionStarting {
+			return &gauge_messages.ProtoExecutionResult{
+				ScreenShot:    [][]byte{[]byte("screenshot1"), []byte("screenshot2")},
+				Failed:        false,
+				ExecutionTime: 10,
+			}
+		}
+		return &gauge_messages.ProtoExecutionResult{}
+	}
+	ei := &gauge_messages.ExecutionInfo{
+		CurrentStep: &gauge_messages.StepInfo{
+			Step: &gauge_messages.ExecuteStepRequest{
+				ActualStepText:  "a simple step",
+				ParsedStepText:  "a simple step",
+				ScenarioFailing: false,
+			},
+			IsFailed: false,
+		},
+	}
+	se := &stepExecutor{runner: r, pluginHandler: h, currentExecutionInfo: ei, stream: 0}
+	step := &gauge.Step{
+		Value:     "a simple step",
+		LineText:  "a simple step",
+		Fragments: []*gauge_messages.Fragment{{FragmentType: gauge_messages.Fragment_Text, Text: "a simple step"}},
+	}
+	protoStep := gauge.ConvertToProtoItem(step).GetStep()
+	protoStep.StepExecutionResult = &gauge_messages.ProtoStepExecutionResult{}
+
+	stepResult := se.executeStep(step, protoStep)
+	beforeStepScreenShots := stepResult.ProtoStep.PreHookScreenshots
+
+	expected := []string{"screenshot1", "screenshot2"}
+
+	if len(beforeStepScreenShots) != len(expected) {
+		t.Errorf("Expected 2 screenshots, got : %d", len(beforeStepScreenShots))
+	}
+
+	for i, e := range expected {
+		if string(beforeStepScreenShots[i]) != e {
+			t.Errorf("Expected `%s` screenshot, got : %s", e, beforeStepScreenShots[i])
+		}
+	}
+}
+
+func TestStepExecutionShouldGetScreenshotsAfterStep(t *testing.T) {
 	r := &mockRunner{}
 	h := &mockPluginHandler{NotifyPluginsfunc: func(m *gauge_messages.Message) {}, GracefullyKillPluginsfunc: func() {}}
 	r.ExecuteAndGetStatusFunc = func(m *gauge_messages.Message) *gauge_messages.ProtoExecutionResult {
@@ -144,17 +192,17 @@ func TestStepExecutionShouldGetScreenshots(t *testing.T) {
 	protoStep.StepExecutionResult = &gauge_messages.ProtoStepExecutionResult{}
 
 	stepResult := se.executeStep(step, protoStep)
-	screenShots := stepResult.ProtoStepExecResult().GetExecutionResult().ScreenShot
+	afterStepScreenShots := stepResult.ProtoStep.PostHookScreenshots
 
 	expected := []string{"screenshot1", "screenshot2"}
 
-	if len(screenShots) != len(expected) {
-		t.Errorf("Expected 1 message, got : %d", len(screenShots))
+	if len(afterStepScreenShots) != len(expected) {
+		t.Errorf("Expected 2 screenshots, got : %d", len(afterStepScreenShots))
 	}
 
 	for i, e := range expected {
-		if string(screenShots[i]) != e {
-			t.Errorf("Expected `%s` message, got : %s", e, screenShots[i])
+		if string(afterStepScreenShots[i]) != e {
+			t.Errorf("Expected `%s` screenshot, got : %s", e, afterStepScreenShots[i])
 		}
 	}
 }

gauge/protoConverters.go

@@ -191,6 +191,8 @@ func ConvertToProtoSuiteResult(suiteResult *result.SuiteResult) *gauge_messages.
 		SpecsSkippedCount: int32(suiteResult.SpecsSkippedCount),
 		PreHookMessages:   suiteResult.PreHookMessages,
 		PostHookMessages:  suiteResult.PostHookMessages,
+		PreHookScreenshots:  suiteResult.PreHookScreenshots,
+		PostHookScreenshots: suiteResult.PostHookScreenshots,
 	}
 	return protoSuiteResult
 }