From f045623fcf17e31ab3a31f43e7cfa51535cfcc30 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fredrik=20L=C3=B6nnblad?= <fredrik+github@lonnblad.com>
Date: Tue, 10 Mar 2020 09:39:59 -0300
Subject: [PATCH] Fixed dependency issue

---
 fmt_pretty.go    |  4 ++--
 suite_context.go | 26 +++++++++++++-------------
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/fmt_pretty.go b/fmt_pretty.go
index 6037d5b..02ad3b9 100644
--- a/fmt_pretty.go
+++ b/fmt_pretty.go
@@ -312,14 +312,14 @@ func (f *pretty) printStep(result *stepResult) {
 	}
 
 	scenarioHeaderLength, maxLength := f.scenarioLengths(result.owner.AstNodeIds[0])
-	stepLength := f.lengthPickleStep(astStep.Keyword, astStep.Text)
+	stepLength := f.lengthPickleStep(astStep.Keyword, result.step.Text)
 
 	firstExecutedScenarioStep := len(f.lastFeature().lastPickleResult().stepResults) == backgroundSteps+1
 	if !astBackgroundStep && firstExecutedScenarioStep {
 		f.printScenarioHeader(astScenario, maxLength-scenarioHeaderLength)
 	}
 
-	text := s(f.indent*2) + result.status.clr()(strings.TrimSpace(astStep.Keyword)) + " " + result.status.clr()(astStep.Text)
+	text := s(f.indent*2) + result.status.clr()(strings.TrimSpace(astStep.Keyword)) + " " + result.status.clr()(result.step.Text)
 	if result.def != nil {
 		text += s(maxLength - stepLength + 1)
 		text += blackb("# " + result.def.definitionID())
diff --git a/suite_context.go b/suite_context.go
index 9544436..73ccb19 100644
--- a/suite_context.go
+++ b/suite_context.go
@@ -107,31 +107,31 @@ func SuiteContext(s *Suite, additionalContextInitializers ...func(suite *Suite))
 
 	s.BeforeStep(c.inject)
 }
-func (s *suiteContext) inject(step *gherkin.Step) {
+
+func (s *suiteContext) inject(step *messages.Pickle_PickleStep) {
 	if !s.allowInjection {
 		return
 	}
 
 	step.Text = injectAll(step.Text)
-	args := step.Argument
-	if args != nil {
-		switch arg := args.(type) {
-		case *gherkin.DataTable:
-			for i := 0; i < len(arg.Rows); i++ {
-				for n, cell := range arg.Rows[i].Cells {
-					arg.Rows[i].Cells[n].Value = injectAll(cell.Value)
-				}
+
+	if table := step.Argument.GetDataTable(); table != nil {
+		for i := 0; i < len(table.Rows); i++ {
+			for n, cell := range table.Rows[i].Cells {
+				table.Rows[i].Cells[n].Value = injectAll(cell.Value)
 			}
-		case *gherkin.DocString:
-			arg.Content = injectAll(arg.Content)
 		}
 	}
+
+	if doc := step.Argument.GetDocString(); doc != nil {
+		doc.Content = injectAll(doc.Content)
+	}
 }
 
-func injectAll(inTo string) string {
+func injectAll(src string) string {
 	re := regexp.MustCompile(`{{[^{}]+}}`)
 	return re.ReplaceAllStringFunc(
-		inTo,
+		src,
 		func(key string) string {
 			injectRegex := regexp.MustCompile(`^{{.+}}$`)
 			if injectRegex.MatchString(key) {