From 2b575519a8c3ea2436104c61609a92d00e661d60 Mon Sep 17 00:00:00 2001
From: Adam Grare <adam@grare.com>
Date: Thu, 20 Jun 2024 10:38:55 -0400
Subject: [PATCH] Fix InputPath/OutputPath for Choice/Succeed states

---
 lib/floe/workflow/states/choice.rb   |  3 ++-
 lib/floe/workflow/states/succeed.rb  | 11 +++++++++-
 spec/workflow/states/succeed_spec.rb | 31 +++++++++++++++++++++++++++-
 3 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/lib/floe/workflow/states/choice.rb b/lib/floe/workflow/states/choice.rb
index b6a5da6da..936497592 100644
--- a/lib/floe/workflow/states/choice.rb
+++ b/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
diff --git a/lib/floe/workflow/states/succeed.rb b/lib/floe/workflow/states/succeed.rb
index 794aff886..741fcdb34 100644
--- a/lib/floe/workflow/states/succeed.rb
+++ b/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
 
diff --git a/spec/workflow/states/succeed_spec.rb b/spec/workflow/states/succeed_spec.rb
index 64fbae1d5..3151c7e41 100644
--- a/spec/workflow/states/succeed_spec.rb
+++ b/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