diff --git a/pkg/yaml/lint.go b/pkg/yaml/lint.go
index cb56c2c9..750f1ff4 100644
--- a/pkg/yaml/lint.go
+++ b/pkg/yaml/lint.go
@@ -50,8 +50,8 @@ func expectImage(c *common.Config) error {
 
 // lint rule that fails when no build commands are defined
 func expectCommand(c *common.Config) error {
-	if c.Build.Config == nil || c.Build.Config["commands"] == nil {
-		return fmt.Errorf("Yaml must define build commands")
+	if c.Setup.Config == nil || c.Setup.Config["commands"] == nil {
+		return fmt.Errorf("Yaml must define build / setup commands")
 	}
 	return nil
 }
diff --git a/pkg/yaml/lint_test.go b/pkg/yaml/lint_test.go
index 0f127435..a7d1afdb 100644
--- a/pkg/yaml/lint_test.go
+++ b/pkg/yaml/lint_test.go
@@ -26,7 +26,7 @@ func Test_Linter(t *testing.T) {
 
 		g.It("Should fail when no commands", func() {
 			c := &common.Config{
-				Build: &common.Step{},
+				Setup: &common.Step{},
 			}
 			g.Assert(expectCommand(c) != nil).IsTrue()
 		})
@@ -46,8 +46,9 @@ func Test_Linter(t *testing.T) {
 			c := &common.Config{}
 			c.Build = &common.Step{}
 			c.Build.Image = "golang"
-			c.Build.Config = map[string]interface{}{}
-			c.Build.Config["commands"] = []string{"go build", "go test"}
+			c.Setup = &common.Step{}
+			c.Setup.Config = map[string]interface{}{}
+			c.Setup.Config["commands"] = []string{"go build", "go test"}
 			c.Clone = &common.Step{}
 			c.Clone.Config = map[string]interface{}{}
 			c.Clone.Config["path"] = "/drone/src/foo/bar"