feat: Add custom description for CloudFormation template (aws#161)

* feat: Add custom description to CFN template

Author: ca-nguyen

src/stepfunctions/workflow/cloudformation.py

@@ -18,11 +18,13 @@
 
 logger = logging.getLogger('stepfunctions')
 
+
 def repr_str(dumper, data):
  if '\n' in data:
  return dumper.represent_scalar(u'tag:yaml.org,2002:str', data, style='|')
  return dumper.org_represent_str(data)
 
+
 yaml.SafeDumper.org_represent_str = yaml.SafeDumper.represent_str
 yaml.add_representer(dict, lambda self, data: yaml.representer.SafeRepresenter.represent_dict(self, data.items()), Dumper=yaml.SafeDumper)
 yaml.add_representer(str, repr_str, Dumper=yaml.SafeDumper)
@@ -42,12 +44,19 @@ def repr_str(dumper, data):
  }
 }
 
-def build_cloudformation_template(workflow):
+
+def build_cloudformation_template(workflow, description=None):
+ """
+ Creates a CloudFormation template from the provided Workflow
+ Args:
+ workflow (Workflow): Step Functions workflow instance
+ description (str, optional): Description of the template. If none provided, the default description will be used: "CloudFormation template for AWS Step Functions - State Machine"
+ """
  logger.warning('To reuse the CloudFormation template in different regions, please make sure to update the region specific AWS resources in the StateMachine definition.')
 
  template = CLOUDFORMATION_BASE_TEMPLATE.copy()
 
- template["Description"] = "CloudFormation template for AWS Step Functions - State Machine"
+ template["Description"] = description if description else "CloudFormation template for AWS Step Functions - State Machine"
  template["Resources"]["StateMachineComponent"]["Properties"]["StateMachineName"] = workflow.name
 
  definition = workflow.definition.to_dict()

src/stepfunctions/workflow/stepfunctions.py

@@ -376,11 +376,13 @@ def render_graph(self, portrait=False):
  widget = WorkflowGraphWidget(self.definition.to_json())
  return widget.show(portrait=portrait)
 
- def get_cloudformation_template(self):
+ def get_cloudformation_template(self, description=None):
  """
  Returns a CloudFormation template that contains only the StateMachine resource. To reuse the CloudFormation template in a different region, please make sure to update the region specific AWS resources (e.g: Lambda ARN, Training Image) in the StateMachine definition.
+ Args:
+ description (str, optional): Description of the template
  """
- return build_cloudformation_template(self)
+ return build_cloudformation_template(self, description)
 
  def __repr__(self):
  return '{}(name={!r}, role={!r}, state_machine_arn={!r})'.format(

tests/unit/test_workflow.py

@@ -56,6 +56,7 @@ def client():
  })
  return sfn
 
+
 @pytest.fixture
 def workflow(client):
  workflow = Workflow(
@@ -67,9 +68,11 @@ def workflow(client):
  workflow.create()
  return workflow
 
+
 def test_workflow_creation(client, workflow):
  assert workflow.state_machine_arn == state_machine_arn
 
+
 def test_workflow_creation_failure_duplicate_state_ids(client):
  improper_definition = steps.Chain([steps.Pass('HelloWorld'), steps.Succeed('HelloWorld')])
  with pytest.raises(ValueError):
@@ -80,6 +83,7 @@ def test_workflow_creation_failure_duplicate_state_ids(client):
  client=client
  )
 
+
 # calling update() before create()
 def test_workflow_update_when_statemachinearn_is_none(client):
  workflow = Workflow(
@@ -92,11 +96,13 @@ def test_workflow_update_when_statemachinearn_is_none(client):
  with pytest.raises(WorkflowNotFound):
  workflow.update(definition=new_definition)
 
+
 # calling update() after create() without arguments
 def test_workflow_update_when_arguments_are_missing(client, workflow):
  with pytest.raises(MissingRequiredParameter):
  workflow.update()
 
+
 # calling update() after create()
 def test_workflow_update(client, workflow):
  client.update_state_machine = MagicMock(return_value={
@@ -106,12 +112,14 @@ def test_workflow_update(client, workflow):
  new_role = 'arn:aws:iam::1234567890:role/service-role/StepFunctionsRoleNew'
  assert workflow.update(definition=new_definition, role=new_role) == state_machine_arn
 
+
 def test_attach_existing_workflow(client):
  workflow = Workflow.attach(state_machine_arn, client)
  assert workflow.name == state_machine_name
  assert workflow.role == role_arn
  assert workflow.state_machine_arn == state_machine_arn
 
+
 def test_workflow_list_executions(client, workflow):
  paginator = client.get_paginator('list_executions')
  paginator.paginate = MagicMock(return_value=[
@@ -140,12 +148,14 @@ def test_workflow_list_executions(client, workflow):
  workflow.state_machine_arn = None
  assert workflow.list_executions() == []
 
+
 def test_workflow_makes_deletion_call(client, workflow):
  client.delete_state_machine = MagicMock(return_value=None)
  workflow.delete()
 
  client.delete_state_machine.assert_called_once_with(stateMachineArn=state_machine_arn)
 
+
 def test_workflow_execute_creation(client, workflow):
  execution = workflow.execute()
  assert execution.workflow.state_machine_arn == state_machine_arn
@@ -164,11 +174,13 @@ def test_workflow_execute_creation(client, workflow):
  input='{}'
  )
 
+
 def test_workflow_execute_when_statemachinearn_is_none(client, workflow):
  workflow.state_machine_arn = None
  with pytest.raises(WorkflowNotFound):
  workflow.execute()
 
+
 def test_execution_makes_describe_call(client, workflow):
  execution = workflow.execute()
 
@@ -177,6 +189,7 @@ def test_execution_makes_describe_call(client, workflow):
 
  client.describe_execution.assert_called_once()
 
+
 def test_execution_makes_stop_call(client, workflow):
  execution = workflow.execute()
 
@@ -194,6 +207,7 @@ def test_execution_makes_stop_call(client, workflow):
  error='Error'
  )
 
+
 def test_execution_list_events(client, workflow):
  paginator = client.get_paginator('get_execution_history')
  paginator.paginate = MagicMock(return_value=[
@@ -229,6 +243,7 @@ def test_execution_list_events(client, workflow):
  }
  )
 
+
 def test_list_workflows(client):
  paginator = client.get_paginator('list_state_machines')
  paginator.paginate = MagicMock(return_value=[
@@ -254,11 +269,14 @@ def test_list_workflows(client):
  }
  )
 
+
 def test_cloudformation_export_with_simple_definition(workflow):
  cfn_template = workflow.get_cloudformation_template()
  cfn_template = yaml.load(cfn_template)
  assert 'StateMachineComponent' in cfn_template['Resources']
  assert workflow.role == cfn_template['Resources']['StateMachineComponent']['Properties']['RoleArn']
+ assert cfn_template['Description'] == "CloudFormation template for AWS Step Functions - State Machine"
+
 
 def test_cloudformation_export_with_sagemaker_execution_role(workflow):
  workflow.definition.to_dict = MagicMock(return_value={
@@ -281,7 +299,8 @@ def test_cloudformation_export_with_sagemaker_execution_role(workflow):
  }
  }
  })
- cfn_template = workflow.get_cloudformation_template()
+ cfn_template = workflow.get_cloudformation_template(description="CloudFormation template with Sagemaker role")
  cfn_template = yaml.load(cfn_template)
  assert json.dumps(workflow.definition.to_dict(), indent=2) == cfn_template['Resources']['StateMachineComponent']['Properties']['DefinitionString']
  assert workflow.role == cfn_template['Resources']['StateMachineComponent']['Properties']['RoleArn']
+ assert cfn_template['Description'] == "CloudFormation template with Sagemaker role"