Make Choice states chainable

CONTRIBUTING.md

@@ -57,8 +57,8 @@ Before sending us a pull request, please ensure that:
 ### Running the Unit Tests
 
 1. Install tox using `pip install tox`
-1. Install test dependencies, including coverage, using `pip install .[test]`
 1. cd into the aws-step-functions-data-science-sdk-python folder: `cd aws-step-functions-data-science-sdk-python` or `cd /environment/aws-step-functions-data-science-sdk-python`
+1. Install test dependencies, including coverage, using `pip install ".[test]"`
 1. Run the following tox command and verify that all code checks and unit tests pass: `tox tests/unit`
 
 You can also run a single test with the following command: `tox -e py36 -- -s -vv <path_to_file><file_name>::<test_function_name>`  
@@ -80,7 +80,7 @@ You should only worry about manually running any new integration tests that you
 
 1. Create a new git branch:
      ```shell
-     git checkout -b my-fix-branch master
+     git checkout -b my-fix-branch main
      ```
 1. Make your changes, **including unit tests** and, if appropriate, integration tests.
    1. Include unit tests when you contribute new features or make bug fixes, as they help to:

src/stepfunctions/steps/states.py

@@ -218,9 +218,21 @@ def next(self, next_step):
         Returns:
             State or Chain: Next state or chain that will be transitioned to.
         """
-        if self.type in ('Choice', 'Succeed', 'Fail'):
+        if self.type in ('Succeed', 'Fail'):
             raise ValueError('Unexpected State instance `{step}`, State type `{state_type}` does not support method `next`.'.format(step=next_step, state_type=self.type))
 
+        # By design, choice states do not have the Next field. Setting default to make it chainable.
+        if self.type is 'Choice':
+            if self.default is not None:
+                logger.warning(
+                    "Chaining Choice Step: Overwriting %s's current default_choice (%s) with %s",
+                    self.state_id,
+                    self.default.state_id,
+                    next_step.state_id
+                )
+            self.default_choice(next_step)
+            return self.default
+
         self.next_step = next_step
         return self.next_step
 

tests/unit/test_steps.py

@@ -12,6 +12,7 @@
 # permissions and limitations under the License.
 from __future__ import absolute_import
 
+import logging
 import pytest
 
 from stepfunctions.exceptions import DuplicateStatesInChain
@@ -328,12 +329,6 @@ def test_append_states_after_terminal_state_will_fail():
         chain.append(Succeed('Succeed'))
         chain.append(Pass('Pass2'))
 
-    with pytest.raises(ValueError):
-        chain = Chain()
-        chain.append(Pass('Pass'))
-        chain.append(Choice('Choice'))
-        chain.append(Pass('Pass2'))
-
 
 def test_chaining_steps():
     s1 = Pass('Step - One')
@@ -372,6 +367,33 @@ def test_chaining_steps():
     assert s1.next_step == s2
     assert s2.next_step == s3
 
+
+def test_chaining_choice(caplog):
+    s1_pass = Pass('Step - One')
+    s2_choice = Choice('Step - Two')
+    s3_pass = Pass('Step - Three')
+
+    with caplog.at_level(logging.WARNING):
+        chain1 = Chain([s1_pass, s2_choice, s3_pass])
+        assert caplog.text == ''  # No warning
+    assert chain1.steps == [s1_pass, s2_choice, s3_pass]
+    assert s1_pass.next_step == s2_choice
+    assert s2_choice.default == s3_pass
+    assert s2_choice.next_step is None  # Choice steps do not have next_step
+    assert s3_pass.next_step is None
+
+    # Chain s2_choice when default_choice is already set will trigger Warning
+    with caplog.at_level(logging.WARNING):
+        Chain([s2_choice, s1_pass])
+        log_message = (
+                "Chaining Choice Step: Overwriting %s's current default_choice (%s) with %s" %
+                (s2_choice.state_id, s3_pass.state_id, s1_pass.state_id)
+        )
+        assert log_message in caplog.text
+    assert s2_choice.default == s1_pass
+    assert s2_choice.next_step is None  # Choice steps do not have next_step
+
+
 def test_catch_fail_for_unsupported_state():
     s1 = Pass('Step - One')