use with_token_refetch during token exchange's activate_shopify_session

Shopify · Apr 10, 2024 · 6c8e444 · 6c8e444
1 parent ba45278
commit 6c8e444
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 6 deletions.
diff --git a/lib/shopify_app.rb b/lib/shopify_app.rb
@@ -40,6 +40,9 @@ def self.use_webpacker?
 
   require "shopify_app/logger"
 
+  # Admin API helpers
+  require "shopify_app/admin_api/with_token_refetch"
+
   # controller concerns
   require "shopify_app/controller_concerns/csrf_protection"
   require "shopify_app/controller_concerns/localization"
@@ -54,9 +57,6 @@ def self.use_webpacker?
   require "shopify_app/controller_concerns/webhook_verification"
   require "shopify_app/controller_concerns/token_exchange"
 
-  # Admin API helpers
-  require "shopify_app/admin_api/with_token_refetch"
-
   # Auth helpers
   require "shopify_app/auth/post_authenticate_tasks"
   require "shopify_app/auth/token_exchange"

diff --git a/lib/shopify_app/controller_concerns/token_exchange.rb b/lib/shopify_app/controller_concerns/token_exchange.rb
@@ -3,14 +3,15 @@
 module ShopifyApp
   module TokenExchange
     extend ActiveSupport::Concern
+    include ShopifyApp::AdminAPI::WithTokenRefetch
 
-    def activate_shopify_session
+    def activate_shopify_session(&block)
       retrieve_session_from_token_exchange if current_shopify_session.blank? || should_exchange_expired_token?
 
       begin
         ShopifyApp::Logger.debug("Activating Shopify session")
         ShopifyAPI::Context.activate_session(current_shopify_session)
-        yield
+        with_token_refetch(current_shopify_session, session_token, &block)
       rescue ShopifyAPI::Errors::HttpResponseError => error
         if error.code == 401
           ShopifyApp::Logger.debug("Admin API returned a 401 Unauthorized error, deleting current access token.")

diff --git a/test/shopify_app/controller_concerns/token_exchange_test.rb b/test/shopify_app/controller_concerns/token_exchange_test.rb
@@ -167,7 +167,19 @@ class TokenExchangeControllerTest < ActionController::TestCase
     end
   end
 
-  test "Deletes existing session and re-raises error when an API 401 error is raised by the action" do
+  test "Wraps action in with_token_refetch" do
+    ShopifyApp::SessionRepository.store_shop_session(@offline_session)
+    ShopifyAPI::Utils::SessionUtils.stubs(:current_session_id).returns(@offline_session_id)
+
+    ApiClass.expects(:perform)
+    @controller.expects(:with_token_refetch).yields
+
+    with_application_test_routes do
+      get :make_api_call, params: { shop: @shop }
+    end
+  end
+
+  test "Deletes existing session and re-raises error when an API 401 error is not fixed by with_token_refetch" do
     ShopifyApp::SessionRepository.store_shop_session(@offline_session)
     ShopifyAPI::Utils::SessionUtils.stubs(:current_session_id).returns(@offline_session_id)
 
@@ -176,6 +188,7 @@ class TokenExchangeControllerTest < ActionController::TestCase
 
     ShopifyAPI::Auth::TokenExchange.expects(:exchange_token).never
     ShopifyApp::SessionRepository.expects(:delete_session).with(@offline_session_id)
+    @controller.stubs(:with_token_refetch).yields
 
     with_application_test_routes do
       response_error = assert_raises(ShopifyAPI::Errors::HttpResponseError) do