lib/floe/workflow/states/wait.rb in floe-0.5.0 vs lib/floe/workflow/states/wait.rb in floe-0.6.0

@@ -4,12 +4,14 @@
 
 module Floe
   class Workflow
     module States
       class Wait < Floe::Workflow::State
-        attr_reader :end, :next, :seconds, :input_path, :output_path
+        include NonTerminalMixin
 
+        attr_reader :end, :input_path, :next, :seconds, :seconds_path, :timestamp, :timestamp_path, :output_path
+
         def initialize(workflow, name, payload)
           super
 
           @next           = payload["Next"]
           @end            = !!payload["End"]
@@ -18,19 +20,25 @@
           @timestamp_path = Path.new(payload["TimestampPath"]) if payload.key?("TimestampPath")
           @seconds_path   = Path.new(payload["SecondsPath"]) if payload.key?("SecondsPath")
 
           @input_path  = Path.new(payload.fetch("InputPath", "$"))
           @output_path = Path.new(payload.fetch("OutputPath", "$"))
+
+          validate_state!
         end
 
         def start(input)
           super
           input = input_path.value(context, input)
 
           context.output     = output_path.value(context, input)
           context.next_state = end? ? nil : @next
-          please_hold(input)
+
+          wait_until!(
+            :seconds => seconds_path ? seconds_path.value(context, input).to_i : seconds,
+            :time    => timestamp_path ? timestamp_path.value(context, input) : timestamp
+          )
         end
 
         def running?
           waiting?
         end
@@ -39,14 +47,11 @@
           @end
         end
 
         private
 
-        def please_hold(input)
-          wait(
-            :seconds => @seconds_path ? @seconds_path.value(context, input).to_i : @seconds,
-            :time    => @timestamp_path ? @timestamp_path.value(context, input) : @timestamp
-          )
+        def validate_state!
+          validate_state_next!
         end
       end
     end
   end
 end