From 6c763c1f0968baddc6696f1b066b163897690f00 Mon Sep 17 00:00:00 2001
From: Puneeth Nanjundaswamy <puneeth.nanjundaswamy@comtravo.com>
Date: Fri, 17 Nov 2017 16:06:23 +0100
Subject: [PATCH] Support Step Function UpdateStateMachine call

---
 aws/resource_aws_sfn_state_machine.go      | 27 ++++++++++++++++++--
 aws/resource_aws_sfn_state_machine_test.go | 29 +++++++++++++++++++---
 2 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/aws/resource_aws_sfn_state_machine.go b/aws/resource_aws_sfn_state_machine.go
index c2eda62bcfb..e59765e48f7 100644
--- a/aws/resource_aws_sfn_state_machine.go
+++ b/aws/resource_aws_sfn_state_machine.go
@@ -1,6 +1,7 @@
 package aws
 
 import (
+	"fmt"
 	"log"
 	"time"
 
@@ -16,6 +17,7 @@ func resourceAwsSfnStateMachine() *schema.Resource {
 	return &schema.Resource{
 		Create: resourceAwsSfnStateMachineCreate,
 		Read:   resourceAwsSfnStateMachineRead,
+		Update: resourceAwsSfnStateMachineUpdate,
 		Delete: resourceAwsSfnStateMachineDelete,
 		Importer: &schema.ResourceImporter{
 			State: schema.ImportStatePassthrough,
@@ -25,7 +27,6 @@ func resourceAwsSfnStateMachine() *schema.Resource {
 			"definition": {
 				Type:         schema.TypeString,
 				Required:     true,
-				ForceNew:     true,
 				ValidateFunc: validateSfnStateMachineDefinition,
 			},
 
@@ -39,7 +40,6 @@ func resourceAwsSfnStateMachine() *schema.Resource {
 			"role_arn": {
 				Type:         schema.TypeString,
 				Required:     true,
-				ForceNew:     true,
 				ValidateFunc: validateArn,
 			},
 
@@ -125,6 +125,29 @@ func resourceAwsSfnStateMachineRead(d *schema.ResourceData, meta interface{}) er
 	return nil
 }
 
+func resourceAwsSfnStateMachineUpdate(d *schema.ResourceData, meta interface{}) error {
+	conn := meta.(*AWSClient).sfnconn
+
+	params := &sfn.UpdateStateMachineInput{
+		StateMachineArn: aws.String(d.Id()),
+		Definition:      aws.String(d.Get("definition").(string)),
+		RoleArn:         aws.String(d.Get("role_arn").(string)),
+	}
+
+	_, err := conn.UpdateStateMachine(params)
+
+	log.Printf("[DEBUG] Updating Step Function State Machine: %#v", params)
+
+	if err != nil {
+		if isAWSErr(err, "StateMachineDoesNotExist", "State Machine Does Not Exist") {
+			return fmt.Errorf("Error updating Step Function State Machine: %s", err)
+		}
+		return err
+	}
+
+	return resourceAwsSfnStateMachineRead(d, meta)
+}
+
 func resourceAwsSfnStateMachineDelete(d *schema.ResourceData, meta interface{}) error {
 	conn := meta.(*AWSClient).sfnconn
 	log.Printf("[DEBUG] Deleting Step Function State Machine: %s", d.Id())
diff --git a/aws/resource_aws_sfn_state_machine_test.go b/aws/resource_aws_sfn_state_machine_test.go
index 408717ac478..7e3da1a0a40 100644
--- a/aws/resource_aws_sfn_state_machine_test.go
+++ b/aws/resource_aws_sfn_state_machine_test.go
@@ -2,6 +2,7 @@ package aws
 
 import (
 	"fmt"
+	"regexp"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -12,7 +13,7 @@ import (
 	"github.com/hashicorp/terraform/terraform"
 )
 
-func TestAccAWSSfnStateMachine_basic(t *testing.T) {
+func TestAccAWSSfnStateMachine_createUpdate(t *testing.T) {
 	name := acctest.RandString(10)
 
 	resource.Test(t, resource.TestCase{
@@ -21,13 +22,25 @@ func TestAccAWSSfnStateMachine_basic(t *testing.T) {
 		CheckDestroy: testAccCheckAWSSfnStateMachineDestroy,
 		Steps: []resource.TestStep{
 			{
-				Config: testAccAWSSfnStateMachineBasicConfig(name),
+				Config: testAccAWSSfnStateMachineConfig(name, 5),
 				Check: resource.ComposeTestCheckFunc(
 					testAccCheckAWSSfnExists("aws_sfn_state_machine.foo"),
 					resource.TestCheckResourceAttr("aws_sfn_state_machine.foo", "status", sfn.StateMachineStatusActive),
 					resource.TestCheckResourceAttrSet("aws_sfn_state_machine.foo", "name"),
 					resource.TestCheckResourceAttrSet("aws_sfn_state_machine.foo", "creation_date"),
 					resource.TestCheckResourceAttrSet("aws_sfn_state_machine.foo", "definition"),
+					resource.TestMatchResourceAttr("aws_sfn_state_machine.foo", "definition", regexp.MustCompile(`.*\"MaxAttempts\": 5.*`)),
+					resource.TestCheckResourceAttrSet("aws_sfn_state_machine.foo", "role_arn"),
+				),
+			},
+			{
+				Config: testAccAWSSfnStateMachineConfig(name, 10),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckAWSSfnExists("aws_sfn_state_machine.foo"),
+					resource.TestCheckResourceAttr("aws_sfn_state_machine.foo", "status", sfn.StateMachineStatusActive),
+					resource.TestCheckResourceAttrSet("aws_sfn_state_machine.foo", "name"),
+					resource.TestCheckResourceAttrSet("aws_sfn_state_machine.foo", "creation_date"),
+					resource.TestMatchResourceAttr("aws_sfn_state_machine.foo", "definition", regexp.MustCompile(`.*\"MaxAttempts\": 10.*`)),
 					resource.TestCheckResourceAttrSet("aws_sfn_state_machine.foo", "role_arn"),
 				),
 			},
@@ -89,7 +102,7 @@ func testAccCheckAWSSfnStateMachineDestroy(s *terraform.State) error {
 	return fmt.Errorf("Default error in Step Function Test")
 }
 
-func testAccAWSSfnStateMachineBasicConfig(rName string) string {
+func testAccAWSSfnStateMachineConfig(rName string, rMaxAttempts int) string {
 	return fmt.Sprintf(`
 data "aws_region" "current" {
   current = true
@@ -190,6 +203,14 @@ resource "aws_sfn_state_machine" "foo" {
     "HelloWorld": {
       "Type": "Task",
       "Resource": "${aws_lambda_function.lambda_function_test.arn}",
+      "Retry": [
+        {
+          "ErrorEquals": ["States.ALL"],
+          "IntervalSeconds": 5,
+          "MaxAttempts": %d,
+          "BackoffRate": 8.0
+        }
+      ],
       "End": true
     }
   }
@@ -197,5 +218,5 @@ resource "aws_sfn_state_machine" "foo" {
 EOF
 }
 
-`, rName, rName, rName, rName, rName, rName)
+`, rName, rName, rName, rName, rName, rName, rMaxAttempts)
 }