chore: Refactor RefreshOauthTokenService to improve readability (#8820)

- Added a trait called microsoft_email for the Channel::Email factory.
- Rewrote the logic to make it simple to understand
- Rewrote the specs for readability

app/services/microsoft/refresh_oauth_token_service.rb

@@ -1,48 +1,39 @@
-# refer: https://gitlab.com/gitlab-org/ruby/gems/gitlab-mail_room/-/blob/master/lib/mail_room/microsoft_graph/connection.rb
+# Refer: https://learn.microsoft.com/en-us/entra/identity-platform/configurable-token-lifetimes
-# refer: https://github.com/microsoftgraph/msgraph-sample-rubyrailsapp/tree/b4a6869fe4a438cde42b161196484a929f1bee46
-# https://learn.microsoft.com/en-us/azure/active-directory/develop/active-directory-configurable-token-lifetimes
 class Microsoft::RefreshOauthTokenService
   pattr_initialize [:channel!]
 
-  # if the token is not expired yet then skip the refresh token step
+  # Additional references: https://gitlab.com/gitlab-org/ruby/gems/gitlab-mail_room/-/blob/master/lib/mail_room/microsoft_graph/connection.rb
   def access_token
-    provider_config = channel.provider_config.with_indifferent_access
+    return provider_config[:access_token] unless access_token_expired?
-    if Time.current.utc >= expires_on(provider_config['expires_on'])
+
-      # Token expired, refresh
+    refreshed_tokens = refresh_tokens
-      new_hash = refresh_tokens
+    refreshed_tokens[:access_token]
-      new_hash[:access_token]
-    else
-      provider_config[:access_token]
-    end
   end
 
-  def expires_on(expiry)
+  def access_token_expired?
-    # we will give it a 5 minute gap for safety
+    expiry = provider_config[:expires_on]
-    expiry.presence ? DateTime.parse(expiry) - 5.minutes : Time.current.utc
+
+    return true if expiry.blank?
+
+    # Adding a 5 minute window to expiry check to avoid any race
+    # conditions during the fetch operation. This would assure that the
+    # tokens are updated when we fetch the emails.
+    Time.current.utc >= DateTime.parse(expiry) - 5.minutes
   end
 
-  # <RefreshTokensSnippet>
+  # Refresh the access tokens using the refresh token
+  # Refer: https://github.com/microsoftgraph/msgraph-sample-rubyrailsapp/tree/b4a6869fe4a438cde42b161196484a929f1bee46
   def refresh_tokens
-    token_hash = channel.provider_config.with_indifferent_access
+    oauth_strategy = build_oauth_strategy
-    oauth_strategy = ::MicrosoftGraphAuth.new(
+    token_service = build_token_service(oauth_strategy)
-      nil, ENV.fetch('AZURE_APP_ID', nil), ENV.fetch('AZURE_APP_SECRET', nil)
-    )
 
-    token_service = OAuth2::AccessToken.new(
-      oauth_strategy.client, token_hash['access_token'],
-      refresh_token: token_hash['refresh_token']
-    )
-
-    # Refresh the tokens
     new_tokens = token_service.refresh!.to_hash.slice(:access_token, :refresh_token, :expires_at)
 
     update_channel_provider_config(new_tokens)
-    channel.provider_config
+    channel.reload.provider_config
   end
-  # </RefreshTokensSnippet>
 
   def update_channel_provider_config(new_tokens)
-    new_tokens = new_tokens.with_indifferent_access
     channel.provider_config = {
       access_token: new_tokens[:access_token],
       refresh_token: new_tokens[:refresh_token],
@@ -50,4 +41,24 @@ class Microsoft::RefreshOauthTokenService
     }
     channel.save!
   end
+
+  private
+
+  def provider_config
+    @provider_config ||= channel.provider_config.with_indifferent_access
+  end
+
+  # Builds the OAuth strategy for Microsoft Graph
+  def build_oauth_strategy
+    ::MicrosoftGraphAuth.new(nil, ENV.fetch('AZURE_APP_ID'), ENV.fetch('AZURE_APP_SECRET'))
+  end
+
+  # Builds the token service using OAuth2
+  def build_token_service(oauth_strategy)
+    OAuth2::AccessToken.new(
+      oauth_strategy.client,
+      provider_config[:access_token],
+      refresh_token: provider_config[:refresh_token]
+    )
+  end
 end

spec/factories/channel/channel_email.rb

@@ -8,5 +8,22 @@ FactoryBot.define do
     after(:create) do |channel_email|
       create(:inbox, channel: channel_email, account: channel_email.account)
     end
+
+    trait :microsoft_email do
+      imap_enabled { true }
+      imap_address { 'outlook.office365.com' }
+      imap_port { 993 }
+      imap_login { 'email@example.com' }
+      imap_password { '' }
+      imap_enable_ssl { true }
+      provider_config do
+        {
+          expires_on: Time.zone.now + 3600,
+          access_token: SecureRandom.hex,
+          refresh_token: SecureRandom.hex
+        }
+      end
+      provider { 'microsoft' }
+    end
   end
 end

spec/jobs/inboxes/fetch_imap_emails_job_spec.rb

@@ -8,11 +8,8 @@ RSpec.describe Inboxes::FetchImapEmailsJob do
     create(:channel_email, imap_enabled: true, imap_address: 'imap.gmail.com', imap_port: 993, imap_login: 'imap@gmail.com',
                            imap_password: 'password', account: account)
   end
-  let(:microsoft_imap_email_channel) do
+
-    create(:channel_email, provider: 'microsoft', imap_enabled: true, imap_address: 'outlook.office365.com',
+  let(:microsoft_imap_email_channel) { create(:channel_email, :microsoft_email) }
-                           imap_port: 993, imap_login: 'imap@outlook.com', imap_password: 'password', account: account,
-                           provider_config: { access_token: 'access_token' })
-  end
   let(:ms_email_inbox) { create(:inbox, channel: microsoft_imap_email_channel, account: account) }
   let!(:conversation) { create(:conversation, inbox: imap_email_channel.inbox, account: account) }
   let(:inbound_mail) { create_inbound_email_from_mail(from: 'testemail@gmail.com', to: 'imap@outlook.com', subject: 'Hello!') }

spec/services/microsoft/refresh_oauth_token_service_spec.rb

@@ -1,51 +1,87 @@
 require 'rails_helper'
 
 RSpec.describe Microsoft::RefreshOauthTokenService do
-  let(:access_token) { SecureRandom.hex }
+  let!(:microsoft_channel) { create(:channel_email, :microsoft_email) }
-  let(:refresh_token) { SecureRandom.hex }
+  let!(:microsoft_channel_with_expired_token) do
-  let(:expires_on) { Time.zone.now + 3600 }
+    create(
-
+      :channel_email, :microsoft_email, provider_config: {
-  let!(:microsoft_email_channel) do
+        expires_on: Time.zone.now - 3600,
-    create(:channel_email, provider_config: { access_token: access_token, refresh_token: refresh_token, expires_on: expires_on })
+        access_token: SecureRandom.hex,
+        refresh_token: SecureRandom.hex
+      }
+    )
   end
-  let(:new_tokens) { { access_token: access_token, refresh_token: refresh_token, expires_at: expires_on.to_i, token_type: 'bearer' } }
 
-  describe '#access_token' do
+  let(:new_tokens) do
-    context 'when token is not expired' do
+    {
-      it 'returns the existing access token' do
+      access_token: SecureRandom.hex,
-        expect(described_class.new(channel: microsoft_email_channel).access_token).to eq(access_token)
+      refresh_token: SecureRandom.hex,
-        expect(microsoft_email_channel.reload.provider_config['refresh_token']).to eq(refresh_token)
+      expires_at: (Time.zone.now + 3600).to_i,
-      end
+      token_type: 'bearer'
+    }
+  end
+
+  context 'when token is not expired' do
+    it 'returns the existing access token' do
+      service = described_class.new(channel: microsoft_channel)
+
+      expect(service.access_token).to eq(microsoft_channel.provider_config['access_token'])
+      expect(microsoft_channel.reload.provider_config['refresh_token']).to eq(microsoft_channel.provider_config['refresh_token'])
+    end
+  end
+
+  describe 'on expired token or invalid expiry' do
+    before do
+      stub_request(:post, 'https://login.microsoftonline.com/common/oauth2/v2.0/token').with(
+        body: { 'grant_type' => 'refresh_token', 'refresh_token' => microsoft_channel_with_expired_token.provider_config['refresh_token'] }
+      ).to_return(status: 200, body: new_tokens.to_json, headers: { 'Content-Type' => 'application/json' })
     end
 
-    context 'when token is expired' do
+    context 'when token is invalid' do
-      let(:expires_on) { 1.minute.from_now }
-
-      before do
-        stub_request(:post, 'https://login.microsoftonline.com/common/oauth2/v2.0/token').with(
-          body: { 'grant_type' => 'refresh_token', 'refresh_token' => refresh_token }
-        ).to_return(status: 200, body: new_tokens.to_json, headers: { 'Content-Type' => 'application/json' })
-      end
-
       it 'fetches new access token and refresh tokens' do
-        microsoft_email_channel.provider_config['expires_on'] = Time.zone.now - 3600
+        with_modified_env AZURE_APP_ID: SecureRandom.uuid, AZURE_APP_SECRET: SecureRandom.hex do
-        microsoft_email_channel.save!
+          provider_config = microsoft_channel_with_expired_token.provider_config
+          service = described_class.new(channel: microsoft_channel_with_expired_token)
+          expect(service.access_token).not_to eq(provider_config['access_token'])
 
-        expect(described_class.new(channel: microsoft_email_channel).access_token).not_to eq(access_token)
+          new_provider_config = microsoft_channel_with_expired_token.reload.provider_config
-        expect(microsoft_email_channel.reload.provider_config['access_token']).to eq(new_tokens[:access_token])
+          expect(new_provider_config['access_token']).to eq(new_tokens[:access_token])
-        expect(microsoft_email_channel.reload.provider_config['refresh_token']).to eq(new_tokens[:refresh_token])
+          expect(new_provider_config['refresh_token']).to eq(new_tokens[:refresh_token])
-        expect(microsoft_email_channel.reload.provider_config['expires_on']).to eq(Time.at(new_tokens[:expires_at]).utc.to_s)
+          expect(new_provider_config['expires_on']).to eq(Time.at(new_tokens[:expires_at]).utc.to_s)
+        end
       end
     end
 
-    context 'when refresh token is not present in provider config and access token is expired' do
+    context 'when expiry time is missing' do
-      it 'throws an error' do
+      it 'fetches new access token and refresh tokens' do
-        microsoft_email_channel.update(provider_config: {
+        with_modified_env AZURE_APP_ID: SecureRandom.uuid, AZURE_APP_SECRET: SecureRandom.hex do
-                                         access_token: access_token,
+          microsoft_channel_with_expired_token.provider_config['expires_on'] = nil
-                                         expires_on: expires_on - 3600
+          microsoft_channel_with_expired_token.save!
-                                       })
+          provider_config = microsoft_channel_with_expired_token.provider_config
+          service = described_class.new(channel: microsoft_channel_with_expired_token)
+          expect(service.access_token).not_to eq(provider_config['access_token'])
+
+          new_provider_config = microsoft_channel_with_expired_token.reload.provider_config
+          expect(new_provider_config['access_token']).to eq(new_tokens[:access_token])
+          expect(new_provider_config['refresh_token']).to eq(new_tokens[:refresh_token])
+          expect(new_provider_config['expires_on']).to eq(Time.at(new_tokens[:expires_at]).utc.to_s)
+        end
+      end
+    end
+  end
+
+  context 'when refresh token is not present in provider config and access token is expired' do
+    it 'throws an error' do
+      with_modified_env AZURE_APP_ID: SecureRandom.uuid, AZURE_APP_SECRET: SecureRandom.hex do
+        microsoft_channel.update(
+          provider_config: {
+            access_token: SecureRandom.hex,
+            expires_on: Time.zone.now - 3600
+          }
+        )
+
         expect do
-          described_class.new(channel: microsoft_email_channel).access_token
+          described_class.new(channel: microsoft_channel).access_token
         end.to raise_error(RuntimeError, 'A refresh_token is not available')
       end
     end