From 17092bb129a39e066e704f52416132640eb042c1 Mon Sep 17 00:00:00 2001
From: Doug Henderson <dhenderson@maestrodev.com>
Date: Thu, 13 Mar 2014 13:13:44 -0700
Subject: [PATCH] Unit test updates + ensure auth retry only done once per
 execute

---
 lib/google/api_client.rb       | 40 ++++++++++++++++-----
 spec/google/api_client_spec.rb | 63 ++++++++++++++++++++++++++++++++--
 2 files changed, 92 insertions(+), 11 deletions(-)

diff --git a/lib/google/api_client.rb b/lib/google/api_client.rb
index c5648081c..913ce4d86 100644
--- a/lib/google/api_client.rb
+++ b/lib/google/api_client.rb
@@ -109,6 +109,7 @@ module Google
       self.key = options[:key]
       self.user_ip = options[:user_ip]
       self.retries = options.fetch(:retries) { 0 }
+      self.expired_auth_retry = options.fetch(:expired_auth_retry) { true }
       @discovery_uris = {}
       @discovery_documents = {}
       @discovered_apis = {}
@@ -239,6 +240,13 @@ module Google
     #  Number of retries
     attr_accessor :retries
 
+    ##
+    # Whether or not an expired auth token should be re-acquired
+    # (and the operation retried) regardless of retries setting
+    # @return [Boolean]
+    #  Auto retry on auth expiry
+    attr_accessor :expired_auth_retry
+
     ##
     # Returns the URI for the directory document.
     #
@@ -593,16 +601,20 @@ module Google
       request.authorization = options[:authorization] || self.authorization unless options[:authenticated] == false
       
       tries = 1 + (options[:retries] || self.retries)
+      attempt = 0
 
       Retriable.retriable :tries => tries, 
                           :on => [TransmissionError],
+                          :on_retry => client_error_handler,
                           :interval => lambda {|attempts| (2 ** attempts) + rand} do
+        attempt += 1
+
         # This 2nd level retriable only catches auth errors, and supports 1 retry, which allows
         # auth to be re-attempted without having to retry all sorts of other failures like
         # NotFound, etc
-        Retriable.retriable :tries => 2, 
+        Retriable.retriable :tries => ((expired_auth_retry || tries > 1) && attempt == 1) ? 2 : 1, 
                             :on => [AuthorizationError],
-                            :on_retry => client_error_handler(request.authorization),
+                            :on_retry => authorization_error_handler(request.authorization),
                             :interval => lambda {|attempts| (2 ** attempts) + rand} do
           result = request.send(connection, true)
 
@@ -671,18 +683,17 @@ module Google
     
 
     ##
-    # Returns on proc for special processing of retries as not all client errors
-    # are recoverable. Only 401s should be retried and only if the credentials
-    # are refreshable
+    # Returns on proc for special processing of retries for authorization errors
+    # Only 401s should be retried and only if the credentials are refreshable
     #
     # @param [#fetch_access_token!] authorization
     #   OAuth 2 credentials
-    # @return [Proc] 
-    def client_error_handler(authorization)  
+    # @return [Proc]
+    def authorization_error_handler(authorization)
       can_refresh = authorization.respond_to?(:refresh_token) && auto_refresh_token 
       Proc.new do |exception, tries|
-        next unless exception.kind_of?(ClientError)
-        if exception.result.status == 401 && can_refresh && tries == 1
+        next unless exception.kind_of?(AuthorizationError)
+        if can_refresh
           begin
             logger.debug("Attempting refresh of access token & retry of request")
             authorization.fetch_access_token!
@@ -694,6 +705,17 @@ module Google
       end
     end
 
+    ##
+    # Returns on proc for special processing of retries as not all client errors
+    # are recoverable. Only 401s should be retried (via authorization_error_handler)
+    #
+    # @return [Proc]
+    def client_error_handler
+      Proc.new do |exception, tries|
+        raise exception if exception.kind_of?(ClientError)
+      end
+    end
+
   end
 
 end
diff --git a/spec/google/api_client_spec.rb b/spec/google/api_client_spec.rb
index 1f4e65054..1cddf8700 100644
--- a/spec/google/api_client_spec.rb
+++ b/spec/google/api_client_spec.rb
@@ -194,7 +194,7 @@ describe Google::APIClient do
       )
     end
 
-    it 'should refresh tokens on 401 tokens' do
+    it 'should refresh tokens on 401 errors' do
       client.authorization.access_token = '12345'
       expect(client.authorization).to receive(:fetch_access_token!)
 
@@ -283,5 +283,64 @@ describe Google::APIClient do
       count.should == 3
     end
 
-  end  
+  end
+
+  describe 'when retries disabled and expired_auth_retry on (default)' do
+    before do
+      client.retries = 0
+    end
+
+    after do
+      @connection.verify
+    end
+
+    it 'should refresh tokens on 401 errors' do
+      client.authorization.access_token = '12345'
+      expect(client.authorization).to receive(:fetch_access_token!)
+
+      @connection = stub_connection do |stub|
+        stub.get('/foo') do |env|
+          [401, {}, '{}']
+        end
+        stub.get('/foo') do |env|
+          [200, {}, '{}']
+        end
+      end
+
+      client.execute(
+        :uri => 'https://www.gogole.com/foo',
+        :connection => @connection
+      )
+    end
+
+  end
+
+  describe 'when retries disabled and expired_auth_retry off' do
+    before do
+      client.retries = 0
+      client.expired_auth_retry = false
+    end
+
+    it 'should not refresh tokens on 401 errors' do
+      client.authorization.access_token = '12345'
+      expect(client.authorization).not_to receive(:fetch_access_token!)
+
+      @connection = stub_connection do |stub|
+        stub.get('/foo') do |env|
+          [401, {}, '{}']
+        end
+        stub.get('/foo') do |env|
+          [200, {}, '{}']
+        end
+      end
+
+      resp = client.execute(
+        :uri => 'https://www.gogole.com/foo',
+        :connection => @connection
+      )
+
+      resp.response.status.should == 401
+    end
+
+  end
 end