lib/calabash-cucumber/wait_helpers.rb in calabash-cucumber-0.9.163.pre6 vs lib/calabash-cucumber/wait_helpers.rb in calabash-cucumber-0.9.163.pre7

@@ -16,11 +16,11 @@
 
       def wait_for(options_or_timeout=
           {:timeout => 10,
            :retry_frequency => 0.2,
            :post_timeout => 0.1,
-           :timeout_message => "Timed out waiting...",
+           :timeout_message => 'Timed out waiting...',
            :screenshot_on_error => true}, &block)
         #note Hash is preferred, number acceptable for backwards compat
         timeout=options_or_timeout
         post_timeout=0.1
         retry_frequency=0.2
@@ -41,11 +41,25 @@
           Timeout::timeout(timeout,WaitError) do
             sleep(retry_frequency) until yield
           end
           sleep(post_timeout) if post_timeout > 0
         rescue WaitError => e
-          handle_error_with_options(e,timeout_message, screenshot_on_error)
+          msg = timeout_message || e
+          if screenshot_on_error
+            sleep(retry_frequency)
+            path  = screenshot
+            res = yield
+            # Validate after taking screenshot
+            if res
+              return res
+            else
+              embed(path)
+              raise msg
+            end
+          else
+            raise msg
+          end
         rescue Exception => e
           handle_error_with_options(e, nil, screenshot_on_error)
         end
       end
 
@@ -66,9 +80,12 @@
         end
       end
 
       #options for wait_for apply
       def wait_for_elements_exist(elements_arr, options={})
+        if elements_arr.is_a?(String)
+          elements_arr = [elements_arr]
+        end
         options[:timeout_message] = options[:timeout_message] || "Timeout waiting for elements: #{elements_arr.join(",")}"
         wait_for(options) do
           elements_arr.all? { |q| element_exists(q) }
         end
       end