From a8e11cd419032fe5b6496f2286ab566eae4293a2 Mon Sep 17 00:00:00 2001
From: Adam Grare <adam@grare.com>
Date: Wed, 22 May 2024 16:27:25 -0400
Subject: [PATCH] Add choice_rule payload validation

---
 lib/floe/workflow/choice_rule/data.rb | 22 ++++++++++++++++++----
 spec/workflow/choice_rule_spec.rb     |  9 +++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/lib/floe/workflow/choice_rule/data.rb b/lib/floe/workflow/choice_rule/data.rb
index fb626fd5a..a21f02be5 100644
--- a/lib/floe/workflow/choice_rule/data.rb
+++ b/lib/floe/workflow/choice_rule/data.rb
@@ -4,6 +4,24 @@ module Floe
   class Workflow
     class ChoiceRule
       class Data < Floe::Workflow::ChoiceRule
+        attr_reader :compare_key
+
+        COMPARE_KEYS = (
+          %w[IsNull IsPresent IsNumeric IsString IsBoolean IsTimestamp StringMatches] +
+          %w[String Numeric Boolean Timestamp].flat_map { |k| ["#{k}Equals", "#{k}EqualsPath"] } +
+          %w[String Numeric Timestamp].flat_map { |k| ["#{k}LessThan", "#{k}LessThanPath", "#{k}GreaterThan", "#{k}GreaterThanPath", "#{k}LessThanEquals", "#{k}LessThanEqualsPath", "#{k}GreaterThanEquals", "#{k}GreaterThanEqualsPath"] }
+        ).freeze
+
+        def initialize(*)
+          super
+
+          compare_keys = payload.keys & COMPARE_KEYS
+          raise Floe::InvalidWorkflowError, "Data-test Expression Choice Rule must have a compare key"        if compare_keys.empty?
+          raise Floe::InvalidWorkflowError, "Data-test Expression Choice Rule must have only one compare key" if compare_keys.size != 1
+
+          @compare_key = compare_keys.first
+        end
+
         def true?(context, input)
           lhs = variable_value(context, input)
           rhs = compare_value(context, input)
@@ -80,10 +98,6 @@ def is_timestamp?(value) # rubocop:disable Naming/PredicateName
           false
         end
 
-        def compare_key
-          @compare_key ||= payload.keys.detect { |key| key.match?(/^(IsNull|IsPresent|IsNumeric|IsString|IsBoolean|IsTimestamp|String|Numeric|Boolean|Timestamp)/) }
-        end
-
         def compare_value(context, input)
           compare_key.end_with?("Path") ? Path.value(payload[compare_key], context, input) : payload[compare_key]
         end
diff --git a/spec/workflow/choice_rule_spec.rb b/spec/workflow/choice_rule_spec.rb
index b08209ada..96683d222 100644
--- a/spec/workflow/choice_rule_spec.rb
+++ b/spec/workflow/choice_rule_spec.rb
@@ -75,6 +75,15 @@
         end
       end
 
+      context "with a missing compare key" do
+        let(:payload) { {"Variable" => "$.foo", "Next" => "FirstMatchState"} }
+        let(:input) { {"foo" => "bar"} }
+
+        it "raises an exception" do
+          expect { subject }.to raise_exception(Floe::InvalidWorkflowError, "Data-test Expression Choice Rule must have a compare key")
+        end
+      end
+
       context "with IsNull" do
         let(:payload) { {"Variable" => "$.foo", "IsNull" => true} }