From 2a4efbf3b9c0e39255c574bfdf9ec57489857fcf Mon Sep 17 00:00:00 2001
From: Sai Cheemalapati <saicheems@gmail.com>
Date: Mon, 18 Sep 2017 12:40:11 -0700
Subject: [PATCH] Allow retry on Signet::AuthorizationError (#611)

---
 lib/google/apis/core/http_command.rb       |  2 +-
 spec/google/apis/core/http_command_spec.rb | 40 +++++++++++++++++++++-
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/lib/google/apis/core/http_command.rb b/lib/google/apis/core/http_command.rb
index 4e544b719..e629b6a93 100644
--- a/lib/google/apis/core/http_command.rb
+++ b/lib/google/apis/core/http_command.rb
@@ -99,7 +99,7 @@ module Google
               # NotFound, etc
               auth_tries = (try == 1 && authorization_refreshable? ? 2 : 1)
               Retriable.retriable tries: auth_tries,
-                                  on: [Google::Apis::AuthorizationError],
+                                  on: [Google::Apis::AuthorizationError, Signet::AuthorizationError],
                                   on_retry: proc { |*| refresh_authorization } do
                 execute_once(client).tap do |result|
                   if block_given?
diff --git a/spec/google/apis/core/http_command_spec.rb b/spec/google/apis/core/http_command_spec.rb
index 964b0e9a0..5750c4acf 100644
--- a/spec/google/apis/core/http_command_spec.rb
+++ b/spec/google/apis/core/http_command_spec.rb
@@ -19,6 +19,44 @@ RSpec.describe Google::Apis::Core::HttpCommand do
   include TestHelpers
   include_context 'HTTP client'
 
+  context('with credentials') do
+    let(:command) do
+      command = Google::Apis::Core::HttpCommand.new(:get, 'https://www.googleapis.com/zoo/animals')
+      command.options.authorization = Signet::OAuth2::Client.new({
+        :token_credential_uri => 'https://accounts.google.com/o/oauth2/token'
+      })
+      command
+    end
+
+    context('with one token error') do
+      before(:example) do
+        stub_request(:post, 'https://accounts.google.com/o/oauth2/token').to_return(
+          { status: [500, 'Server error'] },
+          { status: [200, ''], headers: { 'Content-Type' => 'application/json' }, body: '{"access_token": "foo"}' })
+        stub_request(:get, 'https://www.googleapis.com/zoo/animals').to_return(body: %(Hello world))
+      end
+
+      it 'should send credentials' do
+        result = command.execute(client)
+        expect(a_request(:get, 'https://www.googleapis.com/zoo/animals')
+          .with { |req| req.headers['Authorization'] == 'Bearer foo' }).to have_been_made
+        expect(result).to eql 'Hello world'
+      end
+    end
+
+    context('with two token errors') do
+      before(:example) do
+        stub_request(:post, 'https://accounts.google.com/o/oauth2/token').to_return(
+          { status: [500, 'Server error'] },
+          { status: [401, 'Unauthorized'] })
+      end
+
+      it 'should raise error' do
+        expect { command.execute(client) }.to raise_error(Signet::AuthorizationError)
+      end
+    end
+  end
+
   context('with credentials') do
     let(:command) do
       command = Google::Apis::Core::HttpCommand.new(:get, 'https://www.googleapis.com/zoo/animals')
@@ -29,7 +67,7 @@ RSpec.describe Google::Apis::Core::HttpCommand do
     context('that are refreshable') do
       let(:authorization) do
         calls = 0
-        auth =  object_double(Signet::OAuth2::Client.new)
+        auth = object_double(Signet::OAuth2::Client.new)
         allow(auth).to receive(:apply!) do |header|
           header['Authorization'] = sprintf('Bearer a_token_value_%d', calls)
           calls += 1