diff --git a/features/snippets.feature b/features/snippets.feature
index 5e507f8..e5119f7 100644
--- a/features/snippets.feature
+++ b/features/snippets.feature
@@ -44,11 +44,19 @@ Feature: undefined step snippets
             | col1 | val1 |
             | col2 | val2 |
           Then the response code should be 200 and header "X-Powered-By" should be "godog"
+          And the response body should be:
+          \"\"\"
+          Hello World
+          \"\"\"
       """
     When I run feature suite
     Then the undefined step snippets should be:
       """
-      func iSendRequestToWith(arg1, arg2 string, arg3 *messages.PickleStepArgument_PickleTable) error {
+      func iSendRequestToWith(arg1, arg2 string, arg3 *godog.Table) error {
+              return godog.ErrPending
+      }
+
+      func theResponseBodyShouldBe(arg1 *godog.DocString) error {
               return godog.ErrPending
       }
 
@@ -58,6 +66,7 @@ Feature: undefined step snippets
 
       func InitializeScenario(ctx *godog.ScenarioContext) {
               ctx.Step(`^I send "([^"]*)" request to "([^"]*)" with:$`, iSendRequestToWith)
+              ctx.Step(`^the response body should be:$`, theResponseBodyShouldBe)
               ctx.Step(`^the response code should be (\d+) and header "([^"]*)" should be "([^"]*)"$`, theResponseCodeShouldBeAndHeaderShouldBe)
       }
       """
diff --git a/internal/formatters/undefined_snippets_gen.go b/internal/formatters/undefined_snippets_gen.go
index d47f3bc..792cab9 100644
--- a/internal/formatters/undefined_snippets_gen.go
+++ b/internal/formatters/undefined_snippets_gen.go
@@ -70,11 +70,11 @@ func (s undefinedSnippet) Args() (ret string) {
 
 	if s.argument != nil {
 		if s.argument.GetDocString() != nil {
-			args = append(args, "*messages.PickleStepArgument_PickleDocString")
+			args = append(args, "*godog.DocString")
 		}
 
 		if s.argument.GetDataTable() != nil {
-			args = append(args, "*messages.PickleStepArgument_PickleTable")
+			args = append(args, "*godog.Table")
 		}
 	}