app: retry on throttle

3 files changed, 26 insertions, 2 deletions

diff --git a/lib/hcl/app.rb b/lib/hcl/app.rb
index ac1b992..4185a67 100644
--- a/lib/hcl/app.rb
+++ b/lib/hcl/app.rb
@@ -61,6 +61,10 @@ module HCl
       rescue SocketError => e
         STDERR.puts "Connection failed. (#{e.message})"
         exit 1
+      rescue TimesheetResource::ThrottleFailure => e
+        STDERR.puts "Too many requests, retrying in #{e.retry_after+5} seconds..."
+        sleep e.retry_after+5
+        run
       rescue TimesheetResource::AuthFailure => e
         STDERR.puts "Unable to authenticate: #{e}"
         request_config
diff --git a/lib/hcl/timesheet_resource.rb b/lib/hcl/timesheet_resource.rb
index 69eeb75..d4624ad 100644
--- a/lib/hcl/timesheet_resource.rb
+++ b/lib/hcl/timesheet_resource.rb
@@ -17,6 +17,13 @@ module HCl
   class TimesheetResource
     class Failure < StandardError; end
     class AuthFailure < StandardError; end
+    class ThrottleFailure < StandardError
+      attr_reader :retry_after
+      def initialize response
+        @retry_after = response.headers['Retry-After'].to_i
+        super "Too many requests! Try again in #{@retry_after} seconds."
+      end
+    end
 
     def self.configure opts = nil
       if opts
@@ -68,6 +75,8 @@ module HCl
         response.body
       when Net::HTTPFound
         raise Failure, "Redirected! Perhaps your ssl configuration variable is set incorrectly?"
+      when Net::HTTPServiceUnavailable
+        raise ThrottleFailure, response
       when Net::HTTPUnauthorized
         raise AuthFailure, "Login failed."
       else
diff --git a/test/app_test.rb b/test/app_test.rb
index cbee636..28acba9 100644
--- a/test/app_test.rb
+++ b/test/app_test.rb
@@ -7,11 +7,22 @@ class AppTest < Test::Unit::TestCase
   end
 
   def test_command_show
-    HCl::DayEntry.expects(:all).returns([HCl::DayEntry.new({
+    HCl::DayEntry.expects(:all).returns [HCl::DayEntry.new(
       hours:'2.06',
       notes: 'hi world',
       project: 'App'
-    })])
+    )]
     HCl::App.command 'show'
   end
+
+  def test_command_retry_on_throttle
+    app = HCl::App.new
+    throttled = states('throttled').starts_as(false)
+    app.expects(:show).
+      raises(HCl::TimesheetResource::ThrottleFailure, stub(headers:{'Retry-After' => 42})).
+      then(throttled.is(true))
+    app.expects(:sleep).with(47).when(throttled.is(true))
+    app.expects(:show).when(throttled.is(true))
+    app.process_args('show').run
+  end
 end