Fix InputPath/OutputPath for Choice/Succeed states

lib/floe/workflow/states/choice.rb

@@ -19,7 +19,8 @@ def initialize(workflow, name, payload)
         end
 
         def finish(context)
-          output     = output_path.value(context, context.input)
+          input      = input_path.value(context, context.input)
+          output     = output_path.value(context, input)
           next_state = choices.detect { |choice| choice.true?(context, output) }&.next || default
 
           context.next_state = next_state

lib/floe/workflow/states/succeed.rb

@@ -6,9 +6,18 @@ module States
       class Succeed < Floe::Workflow::State
         attr_reader :input_path, :output_path
 
+        def initialize(workflow, name, payload)
+          super
+
+          @input_path  = Path.new(payload.fetch("InputPath", "$"))
+          @output_path = Path.new(payload.fetch("OutputPath", "$"))
+        end
+
         def finish(context)
+          input              = input_path.value(context, context.input)
+          context.output     = output_path.value(context, input)
           context.next_state = nil
-          context.output     = context.input
+
           super
         end
 

spec/workflow/states/succeed_spec.rb

@@ -2,7 +2,8 @@
   let(:input)    { {} }
   let(:ctx)      { Floe::Workflow::Context.new(:input => input) }
   let(:state)    { workflow.start_workflow.current_state }
-  let(:workflow) { make_workflow(ctx, {"SuccessState" => {"Type" => "Succeed"}}) }
+  let(:payload)  { {"SuccessState" => {"Type" => "Succeed"}} }
+  let(:workflow) { make_workflow(ctx, payload) }
 
   it "#end?" do
     expect(state.end?).to be true
@@ -13,5 +14,33 @@
       state.run_nonblock!(ctx)
       expect(ctx.next_state).to eq(nil)
     end
+
+    context "with input" do
+      let(:input) { {"color" => "red"} }
+
+      it "sets output to input" do
+        state.run_nonblock!(ctx)
+        expect(ctx.output).to eq(input)
+      end
+
+      context "with InputPath" do
+        let(:payload) { {"SuccessState" => {"Type" => "Succeed", "InputPath" => "$.color"}} }
+
+        it "sets the output to the selected input path" do
+          state.run_nonblock!(ctx)
+          expect(ctx.output).to eq(input["color"])
+        end
+      end
+
+      context "with OutputPath" do
+        let(:input)   { {"color" => "red", "garbage" => nil} }
+        let(:payload) { {"SuccessState" => {"Type" => "Succeed", "OutputPath" => "$.color"}} }
+
+        it "sets the output to the selected input path" do
+          state.run_nonblock!(ctx)
+          expect(ctx.output).to eq(input["color"])
+        end
+      end
+    end
   end
 end