From c370f5511262bc754c61e56950b84a991ffd0cdf Mon Sep 17 00:00:00 2001
From: Ami Mahloof <130996527+ami-descope@users.noreply.github.com>
Date: Fri, 6 Sep 2024 01:29:33 -0400
Subject: [PATCH] Handle token in cookie, fix refresh session bug. (#76)

* Handle token in cookie, fix refresh session bug.

* fix order

* fix tests

* fix lint

* fix tests

* move managemnt key check to management tests
---
 lib/descope/api/v1/auth.rb                    | 34 ++++++++++---
 lib/descope/api/v1/auth/enchantedlink.rb      |  4 +-
 lib/descope/api/v1/auth/magiclink.rb          |  4 +-
 lib/descope/api/v1/auth/otp.rb                |  4 +-
 lib/descope/api/v1/auth/password.rb           |  8 ++-
 lib/descope/api/v1/auth/totp.rb               |  4 +-
 lib/descope/api/v1/session.rb                 |  6 ++-
 lib/descope/mixins/http.rb                    | 19 +++++--
 .../lib.descope/api/v1/auth/session_spec.rb   | 49 +++++++++++++++++++
 .../api/v1/management/access_key_spec.rb      |  3 ++
 .../api/v1/management/audit_spec.rb           |  5 +-
 .../api/v1/management/authz_spec.rb           |  2 +
 .../api/v1/management/flow_spec.rb            |  2 +
 .../api/v1/management/permissions_spec.rb     |  2 +
 .../api/v1/management/project_spec.rb         |  2 +
 .../api/v1/management/roles_spec.rb           |  2 +
 .../api/v1/management/user_spec.rb            | 10 ++--
 .../api/v1/auth/enchantedlink_spec.rb         | 13 ++++-
 spec/lib.descope/api/v1/auth/password_spec.rb | 11 ++++-
 spec/lib.descope/api/v1/auth_spec.rb          | 22 +++++++--
 spec/lib.descope/api/v1/session_spec.rb       | 32 +++++++++---
 spec/support/client_config.rb                 |  1 -
 22 files changed, 198 insertions(+), 41 deletions(-)
 create mode 100644 spec/integration/lib.descope/api/v1/auth/session_spec.rb

diff --git a/lib/descope/api/v1/auth.rb b/lib/descope/api/v1/auth.rb
index fc7f6c1..8754ad1 100644
--- a/lib/descope/api/v1/auth.rb
+++ b/lib/descope/api/v1/auth.rb
@@ -34,7 +34,6 @@ def generate_jwt_response(response_body: nil, refresh_cookie: nil, audience: nil
           end
 
           jwt_response = generate_auth_info(response_body, refresh_cookie, true, audience)
-          @logger.debug "jwt_response: #{jwt_response}"
           jwt_response['user'] = response_body.key?('user') ? response_body['user'] : {}
           jwt_response['firstSeen'] = response_body.key?('firstSeen') ? response_body['firstSeen'] : true
 
@@ -62,6 +61,8 @@ def select_tenant(tenant_id: nil, refresh_token: nil)
           validate_refresh_token_not_nil(refresh_token)
           res = post(SELECT_TENANT_PATH, { tenantId: tenant_id }, {}, refresh_token)
           @logger.debug "select_tenant response: #{res}"
+          cookies = res.fetch('cookies')
+          generate_jwt_response(response_body: res, refresh_cookie: cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil))
           generate_jwt_response(
             response_body: res,
             refresh_cookie: res['refreshJwt']
@@ -231,24 +232,41 @@ def validate_token(token, _audience = nil)
         private
 
         def generate_auth_info(response_body, refresh_token, user_jwt, audience = nil)
-          @logger.debug "generating auth info: #{response_body}, #{refresh_token}, #{user_jwt}, #{audience}"
+          @logger.debug "generating auth info: response_body: #{response_body}, refresh_token: #{refresh_token}, user_jwt: #{user_jwt}, audience: #{audience}"
           jwt_response = {}
 
           # validate the session token if sessionJwt is not empty
           st_jwt = response_body.fetch('sessionJwt', '')
           unless st_jwt.empty?
-            @logger.debug "validating session token with refresh_token: #{refresh_token}" if st_jwt
-            jwt_response[SESSION_TOKEN_NAME] = validate_token(st_jwt, audience) if st_jwt
+            @logger.debug 'found sessionJwt in response body, adding to jwt_response'
+            jwt_response[SESSION_TOKEN_NAME] = validate_token(st_jwt, audience)
           end
 
           # validate refresh token if refresh_token was passed or if refreshJwt is not empty
           rt_jwt = response_body.fetch('refreshJwt', '')
 
-          if !refresh_token.nil? || !refresh_token.to_s.empty?
-            @logger.debug "validating refresh token: #{refresh_token}" if refresh_token
-            jwt_response[REFRESH_SESSION_TOKEN_NAME] = validate_token(refresh_token, audience)
-          elsif !rt_jwt.empty?
+          if !rt_jwt.empty?
+            @logger.debug 'found refreshJwt in response body, adding to jwt_response'
+            @logger.debug 'validating refreshJwt token...'
             jwt_response[REFRESH_SESSION_TOKEN_NAME] = validate_token(rt_jwt, audience)
+          elsif refresh_token && !refresh_token.empty?
+            # if refresh_token is in response body (local storage)
+            @logger.debug 'refreshJwt is empty, but refresh_token was passed, adding to jwt_response'
+            @logger.debug 'validating passed-in refresh token...'
+            jwt_response[REFRESH_SESSION_TOKEN_NAME] = validate_token(refresh_token, audience)
+          else
+            cookies = response_body.fetch('cookies', {})
+            # else if refresh token is in response cookie
+            cookies.each do |cookie_name, cookie_value|
+              if cookie_name == REFRESH_SESSION_COOKIE_NAME
+                jwt_response[REFRESH_SESSION_TOKEN_NAME] = validate_token(cookie_value, audience)
+              end
+            end
+          end
+
+          if jwt_response[REFRESH_SESSION_TOKEN_NAME].nil?
+            @logger.debug "Error: Could not find refreshJwt in response body: #{response_body} / cookies: #{cookies} / passed in refresh_token ->#{refresh_token}<-"
+            raise Descope::AuthException.new('Could not find refreshJwt in response body / cookies / passed in refresh_token', code: 500)
           end
 
           jwt_response = adjust_properties(jwt_response, user_jwt)
diff --git a/lib/descope/api/v1/auth/enchantedlink.rb b/lib/descope/api/v1/auth/enchantedlink.rb
index eba9c70..8d77246 100644
--- a/lib/descope/api/v1/auth/enchantedlink.rb
+++ b/lib/descope/api/v1/auth/enchantedlink.rb
@@ -67,7 +67,9 @@ def enchanted_link_verify_token(token = nil)
           def enchanted_link_get_session(pending_ref = nil)
             # @see https://docs.descope.com/api/openapi/enchantedlink/operation/GetEnchantedLinkSession/
             res = post(GET_SESSION_ENCHANTEDLINK_AUTH_PATH, { pendingRef: pending_ref })
-            generate_jwt_response(response_body: res, refresh_cookie: res['refreshJwt'])
+            cookies = res.fetch(COOKIE_DATA_NAME, {})
+            refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+            generate_jwt_response(response_body: res, refresh_cookie:)
           end
 
           private
diff --git a/lib/descope/api/v1/auth/magiclink.rb b/lib/descope/api/v1/auth/magiclink.rb
index 13ce123..4b1057e 100644
--- a/lib/descope/api/v1/auth/magiclink.rb
+++ b/lib/descope/api/v1/auth/magiclink.rb
@@ -42,7 +42,9 @@ def magiclink_sign_up_or_in(method: nil, login_id: nil, uri: nil, login_options:
           def magiclink_verify_token(token = nil)
             validate_token_not_empty(token)
             res = post(VERIFY_MAGICLINK_AUTH_PATH, { token: })
-            generate_jwt_response(response_body: res)
+            cookies = res.fetch(COOKIE_DATA_NAME, {})
+            refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+            generate_jwt_response(response_body: res, refresh_cookie:)
           end
 
           def magiclink_update_user_email(login_id: nil, email: nil, uri: nil, add_to_login_ids: nil, on_merge_use_existing: nil, provider_id: nil, template_id: nil, template_options: nil, refresh_token: nil)
diff --git a/lib/descope/api/v1/auth/otp.rb b/lib/descope/api/v1/auth/otp.rb
index 2aea306..b72aa0b 100644
--- a/lib/descope/api/v1/auth/otp.rb
+++ b/lib/descope/api/v1/auth/otp.rb
@@ -63,7 +63,9 @@ def otp_verify_code(method: nil, login_id: nil, code: nil)
               code:
             }
             res = post(uri, request_params)
-            generate_jwt_response(response_body: res, refresh_cookie: res.fetch('refreshJwt', {}))
+            cookies = res.fetch(COOKIE_DATA_NAME, {})
+            refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+            generate_jwt_response(response_body: res, refresh_cookie:)
           end
 
           def otp_update_user_email(login_id: nil, email: nil, refresh_token: nil, add_to_login_ids: false,
diff --git a/lib/descope/api/v1/auth/password.rb b/lib/descope/api/v1/auth/password.rb
index 54cf984..1e6888f 100644
--- a/lib/descope/api/v1/auth/password.rb
+++ b/lib/descope/api/v1/auth/password.rb
@@ -22,7 +22,9 @@ def password_sign_up(login_id: nil, password: nil, user: nil)
 
             request_params[:user] = password_user_compose_update_body(**user) unless user.nil?
             res = post(SIGN_UP_PASSWORD_PATH, request_params)
-            generate_jwt_response(response_body: res)
+            cookies = res.fetch(COOKIE_DATA_NAME, {})
+            refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+            generate_jwt_response(response_body: res, refresh_cookie:)
           end
 
           def password_sign_in(login_id: nil, password: nil, sso_app_id: nil)
@@ -38,7 +40,9 @@ def password_sign_in(login_id: nil, password: nil, sso_app_id: nil)
               ssoAppId: sso_app_id
             }
             res = post(SIGN_IN_PASSWORD_PATH, request_params)
-            generate_jwt_response(response_body: res)
+            cookies = res.fetch(COOKIE_DATA_NAME, {})
+            refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+            generate_jwt_response(response_body: res, refresh_cookie:)
           end
 
           def password_replace(login_id: nil, old_password: nil, new_password: nil)
diff --git a/lib/descope/api/v1/auth/totp.rb b/lib/descope/api/v1/auth/totp.rb
index 1429f8c..95100e1 100644
--- a/lib/descope/api/v1/auth/totp.rb
+++ b/lib/descope/api/v1/auth/totp.rb
@@ -17,7 +17,9 @@ def totp_sign_in_code(login_id: nil, login_options: nil, code: nil)
             uri = VERIFY_TOTP_PATH
             body = totp_compose_signin_body(login_id, code, login_options)
             res = post(uri, body, {}, nil)
-            generate_jwt_response(response_body: res, refresh_cookie: res.fetch('refreshJwt', {}))
+            cookies = res.fetch(COOKIE_DATA_NAME, {})
+            refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+            generate_jwt_response(response_body: res, refresh_cookie:)
           end
 
           def totp_sign_up(login_id: nil, user: nil, sso_app_id: nil)
diff --git a/lib/descope/api/v1/session.rb b/lib/descope/api/v1/session.rb
index 5a2d0a1..2ddebff 100644
--- a/lib/descope/api/v1/session.rb
+++ b/lib/descope/api/v1/session.rb
@@ -20,11 +20,13 @@ def refresh_session(refresh_token: nil, audience: nil)
           #  [amr, drn, exp, iss, rexp, sub, jwt] in the top level of the response dict, please use
           #  them from the sessionToken key instead, as these claims will soon be deprecated from the top level
           #  of the response dict.
-
+          #  Make sure you set Enable refresh token rotation in the Project Settings before using this.
           validate_refresh_token_not_nil(refresh_token)
           validate_token(refresh_token, audience)
           res = post(REFRESH_TOKEN_PATH, {}, {}, refresh_token)
-          generate_jwt_response(response_body: res, refresh_cookie: refresh_token, audience:)
+          cookies = res.fetch(COOKIE_DATA_NAME, {})
+          refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+          generate_jwt_response(response_body: res, refresh_cookie:, audience:)
         end
 
         def me(refresh_token = nil)
diff --git a/lib/descope/mixins/http.rb b/lib/descope/mixins/http.rb
index d013d0b..0765f77 100644
--- a/lib/descope/mixins/http.rb
+++ b/lib/descope/mixins/http.rb
@@ -1,5 +1,5 @@
 # frozen_string_literal: true
-require "addressable/uri"
+require 'addressable/uri'
 require 'retryable'
 require_relative '../exception'
 
@@ -44,9 +44,17 @@ def retry_options
         }
       end
 
-      def safe_parse_json(body)
+      def safe_parse_json(body, cookies: {})
         @logger.debug "response => #{JSON.parse(body.to_s)}"
-        JSON.parse(body.to_s)
+        res = JSON.parse(body.to_s)
+
+        # Handle DSR cookie in response.
+        if cookies.key?(REFRESH_SESSION_COOKIE_NAME)
+          res['cookies'] = {}
+          res['cookies'][REFRESH_SESSION_COOKIE_NAME] = cookies[REFRESH_SESSION_COOKIE_NAME]
+        end
+
+        res
       rescue JSON::ParserError
         body
       end
@@ -94,11 +102,12 @@ def request(method, uri, body = {}, extra_headers = {})
                    call(method, encode_uri(uri), timeout, @headers, body.to_json)
                  end
 
-        raise Descope::Unsupported.new("No response from server", code: 400) unless result && result.respond_to?(:code)
+        raise Descope::Unsupported.new('No response from server', code: 400) unless result.respond_to?(:code)
 
         @logger.info("API Request: [#{method}] #{uri} - Response Code: #{result.code}")
+
         case result.code
-        when 200...226 then safe_parse_json(result.body)
+        when 200...226 then safe_parse_json(result.body, cookies: result.cookies)
         when 400       then raise Descope::BadRequest.new(result.body, code: result.code, headers: result.headers)
         when 401       then raise Descope::Unauthorized.new(result.body, code: result.code, headers: result.headers)
         when 403       then raise Descope::AccessDenied.new(result.body, code: result.code, headers: result.headers)
diff --git a/spec/integration/lib.descope/api/v1/auth/session_spec.rb b/spec/integration/lib.descope/api/v1/auth/session_spec.rb
new file mode 100644
index 0000000..81c7275
--- /dev/null
+++ b/spec/integration/lib.descope/api/v1/auth/session_spec.rb
@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Descope::Api::V1::Session do
+  before(:all) do
+    @client = DescopeClient.new(Configuration.config)
+  end
+
+  after(:all) do
+    @client.logger.info('Cleaning up test users...')
+    all_users = @client.search_all_users
+    all_users['users'].each do |user|
+      if user['middleName'] == 'Ruby SDK User'
+        @client.logger.info("Deleting ruby spec test user #{user['loginIds'][0]}")
+        @client.delete_user(user['loginIds'][0])
+      end
+    end
+  end
+
+  context 'test session methods' do
+    it 'should refresh session with refresh token' do
+      @password = SpecUtils.generate_password
+      user = build(:user)
+
+      @client.logger.info('1. Sign up with password')
+      res = @client.password_sign_up(login_id: user[:login_id], password: @password, user:)
+      @client.logger.info("sign up with password res: #{res}")
+      original_refresh_token = res[REFRESH_SESSION_TOKEN_NAME]['jwt']
+
+      @client.logger.info('2. Sign in with password')
+      login_res = @client.password_sign_in(login_id: user[:login_id], password: @password)
+      @client.logger.info("sign_in res: #{login_res}")
+
+      @client.logger.info('3. sleep 1 second before calling refresh_session')
+      sleep(1)
+
+      @client.logger.info('4. Refresh session')
+      refresh_session_res = @client.refresh_session(refresh_token: login_res[REFRESH_SESSION_TOKEN_NAME]['jwt'])
+      @client.logger.info("refresh_session_res: #{refresh_session_res}")
+
+      new_refresh_token = refresh_session_res[REFRESH_SESSION_TOKEN_NAME]['jwt']
+      @client.logger.info("new_refresh_token: #{new_refresh_token}")
+
+      @client.logger.info('5. Check new refresh token is not the same as the original one')
+      expect(original_refresh_token).not_to eq(new_refresh_token)
+    end
+  end
+end
diff --git a/spec/integration/lib.descope/api/v1/management/access_key_spec.rb b/spec/integration/lib.descope/api/v1/management/access_key_spec.rb
index d3d9ce7..795776d 100644
--- a/spec/integration/lib.descope/api/v1/management/access_key_spec.rb
+++ b/spec/integration/lib.descope/api/v1/management/access_key_spec.rb
@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::AccessKey do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
   end
 
@@ -28,6 +30,7 @@
       @tenant_id = @client.create_tenant(name: 'some-new-tenant')['id']
       @client.logger.info('creating access key')
       @access_key = @client.create_access_key(name: @key_name, key_tenants: [{ tenant_id: @tenant_id }])
+      @client.logger.info("waiting for access key #{@access_key['key']['id']} to be active 60 seconds")
       sleep 60
     end
 
diff --git a/spec/integration/lib.descope/api/v1/management/audit_spec.rb b/spec/integration/lib.descope/api/v1/management/audit_spec.rb
index a44dde2..c9afd3a 100644
--- a/spec/integration/lib.descope/api/v1/management/audit_spec.rb
+++ b/spec/integration/lib.descope/api/v1/management/audit_spec.rb
@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::Audit do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
     @client.logger.info('Deleting all tenants for Ruby SDK...')
     @client.search_all_tenants(names: ['Ruby-SDK-test'])['tenants'].each do |tenant|
@@ -39,7 +41,7 @@
     created_user = @client.create_user(**user)['user']
 
     expect do
-      res = @client.audit_create_event(
+      @client.audit_create_event(
         user_id: created_user['loginId'],
         action: 'pencil.created',
         type: 'info',
@@ -47,7 +49,6 @@
         actor_id: created_user['loginIds'][0],
         data: { 'key' => 'value' }
       )
-      expect(res).to eq({})
     end.not_to raise_error
   end
 end
diff --git a/spec/integration/lib.descope/api/v1/management/authz_spec.rb b/spec/integration/lib.descope/api/v1/management/authz_spec.rb
index 2af6eb9..4dea5a0 100644
--- a/spec/integration/lib.descope/api/v1/management/authz_spec.rb
+++ b/spec/integration/lib.descope/api/v1/management/authz_spec.rb
@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::Authz do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
     puts 'authz schema delete'
   end
diff --git a/spec/integration/lib.descope/api/v1/management/flow_spec.rb b/spec/integration/lib.descope/api/v1/management/flow_spec.rb
index b8ce624..2e253e1 100644
--- a/spec/integration/lib.descope/api/v1/management/flow_spec.rb
+++ b/spec/integration/lib.descope/api/v1/management/flow_spec.rb
@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::Flow do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
   end
 
diff --git a/spec/integration/lib.descope/api/v1/management/permissions_spec.rb b/spec/integration/lib.descope/api/v1/management/permissions_spec.rb
index bd0b777..d1be727 100644
--- a/spec/integration/lib.descope/api/v1/management/permissions_spec.rb
+++ b/spec/integration/lib.descope/api/v1/management/permissions_spec.rb
@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::Permission do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
     @client.load_all_permissions['permissions'].each do |perm|
       if perm['description'] == 'Ruby SDK'
diff --git a/spec/integration/lib.descope/api/v1/management/project_spec.rb b/spec/integration/lib.descope/api/v1/management/project_spec.rb
index 7a60391..d441c93 100644
--- a/spec/integration/lib.descope/api/v1/management/project_spec.rb
+++ b/spec/integration/lib.descope/api/v1/management/project_spec.rb
@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::Project do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
     @export_output = @client.export_project
   end
diff --git a/spec/integration/lib.descope/api/v1/management/roles_spec.rb b/spec/integration/lib.descope/api/v1/management/roles_spec.rb
index 0819b2f..ea0ed03 100644
--- a/spec/integration/lib.descope/api/v1/management/roles_spec.rb
+++ b/spec/integration/lib.descope/api/v1/management/roles_spec.rb
@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::Role do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
     @client.logger.info('Staring cleanup before tests...')
     @client.logger.info('Deleting all permissions for Ruby SDK...')
diff --git a/spec/integration/lib.descope/api/v1/management/user_spec.rb b/spec/integration/lib.descope/api/v1/management/user_spec.rb
index 10692ec..0c3f841 100644
--- a/spec/integration/lib.descope/api/v1/management/user_spec.rb
+++ b/spec/integration/lib.descope/api/v1/management/user_spec.rb
@@ -4,10 +4,14 @@
 
 describe Descope::Api::V1::Management::User do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
+    @client = DescopeClient.new(Configuration.config)
+
     @password = SpecUtils.generate_password
     @new_password = SpecUtils.generate_password
     @user = build(:user)
-    @client = DescopeClient.new(Configuration.config)
+
     include Descope::Mixins::Common::DeliveryMethod
   end
 
@@ -226,8 +230,8 @@
       new_password = SpecUtils.generate_password
       @client.set_password(login_id: user['loginIds'][0], password: new_password)
       @client.password_sign_in(login_id: user['loginIds'][0], password:)
-    rescue Descope::Unauthorized => e
-      expect(e.message).to match(/"errorDescription":"Invalid signin credentials"/)
+    rescue Descope::ServerError => e
+      expect(e.message).to match(/"Password expired"/)
     end
   end
 
diff --git a/spec/lib.descope/api/v1/auth/enchantedlink_spec.rb b/spec/lib.descope/api/v1/auth/enchantedlink_spec.rb
index 4328e59..7a848c7 100644
--- a/spec/lib.descope/api/v1/auth/enchantedlink_spec.rb
+++ b/spec/lib.descope/api/v1/auth/enchantedlink_spec.rb
@@ -51,7 +51,7 @@
 
     it 'is expected to validate refresh token and not raise an error with refresh token and valid login options' do
       expect do
-        @instance.send(:validate_refresh_token_provided, { mfa: true, stepup: true },  'some-token')
+        @instance.send(:validate_refresh_token_provided, { mfa: true, stepup: true }, 'some-token')
       end.not_to raise_error
     end
 
@@ -148,7 +148,16 @@
     end
 
     it 'is expected to get session by pending ref with enchanted link' do
-      jwt_response = { 'fake': 'response' }
+      jwt_response = {
+        'sessionJwt' => 'fake_session_jwt',
+        'refreshJwt' => 'fake_refresh_jwt',
+        'cookies' => {
+          'refresh_token' => 'fake_refresh_cookie'
+        }
+      }
+      allow(@instance).to receive(:post).with(
+        GET_SESSION_ENCHANTEDLINK_AUTH_PATH, { pendingRef: 'pendingRef' }
+      ).and_return(jwt_response)
       allow(@instance).to receive(:generate_jwt_response).and_return(jwt_response)
 
       expect do
diff --git a/spec/lib.descope/api/v1/auth/password_spec.rb b/spec/lib.descope/api/v1/auth/password_spec.rb
index d976253..abb2ced 100644
--- a/spec/lib.descope/api/v1/auth/password_spec.rb
+++ b/spec/lib.descope/api/v1/auth/password_spec.rb
@@ -43,9 +43,18 @@
     end
 
     it 'is expected to sign in with password' do
+      response_body = {
+        'sessionJwt' => 'fake_session_jwt',
+        'refreshJwt' => 'fake_refresh_jwt',
+        'cookies' => {
+          'refresh_token' => 'fake_refresh_cookie'
+        }
+      }
+
       expect(@instance).to receive(:post).with(
         SIGN_IN_PASSWORD_PATH, { loginId: 'test', password: 's3cr3t', ssoAppId: nil }
-      )
+      ).and_return(response_body)
+
       # stub the jwt_get_unverified_header method to return the kid of the public key created above
       allow(@instance).to receive(:generate_jwt_response).and_return({})
       expect { @instance.password_sign_in(login_id: 'test', password: 's3cr3t') }.not_to raise_error
diff --git a/spec/lib.descope/api/v1/auth_spec.rb b/spec/lib.descope/api/v1/auth_spec.rb
index f67afef..4c38794 100644
--- a/spec/lib.descope/api/v1/auth_spec.rb
+++ b/spec/lib.descope/api/v1/auth_spec.rb
@@ -249,7 +249,13 @@
     end
 
     it 'is expected to select tenant' do
-      jwt_response = { 'fake': 'response' }
+      jwt_response = {
+        'sessionJwt' => 'fake_session_jwt',
+        'refreshJwt' => 'fake_refresh_jwt',
+        'cookies' => {
+          'refresh_token' => 'fake_refresh_cookie'
+        }
+      }
 
       expect(@instance).to receive(:post).with(
         SELECT_TENANT_PATH, { tenantId: 'tenant123' }, {}, 'refresh-token'
@@ -390,20 +396,26 @@
     end
 
     it 'is expected to successfully exchange access key without login_options' do
-      jwt_response = { 'fake': 'response' }
+      jwt_response = {
+        'sessionJwt' => 'fake_session_jwt',
+        'refreshJwt' => 'fake_refresh_jwt'
+      }
       access_key = 'abc'
 
       expect(@instance).to receive(:post).with(
         EXCHANGE_AUTH_ACCESS_KEY_PATH, { loginOptions: {}, audience: 'IT' }, {}, access_key
       ).and_return(jwt_response)
 
-      allow(@instance).to receive(:generate_jwt_response).and_return(jwt_response)
+      allow(@instance).to receive(:generate_auth_info).and_return(jwt_response)
 
       expect { @instance.exchange_access_key(access_key:, audience: 'IT') }.not_to raise_error
     end
 
     it 'is expected to successfully exchange access key with login_options' do
-      jwt_response = { 'fake': 'response' }
+      jwt_response = {
+        'sessionJwt' => 'fake_session_jwt',
+        'refreshJwt' => 'fake_refresh_jwt'
+      }
       access_key = 'abc'
 
       expect(@instance).to receive(:post).with(
@@ -413,7 +425,7 @@
         access_key
       ).and_return(jwt_response)
 
-      allow(@instance).to receive(:generate_jwt_response).and_return(jwt_response)
+      allow(@instance).to receive(:generate_auth_info).and_return(jwt_response)
 
       expect { @instance.exchange_access_key(access_key:, login_options: { customClaims: { k1: 'v1' } }, audience: 'IT') }.not_to raise_error
     end
diff --git a/spec/lib.descope/api/v1/session_spec.rb b/spec/lib.descope/api/v1/session_spec.rb
index 2e1a2fa..87aeb9a 100644
--- a/spec/lib.descope/api/v1/session_spec.rb
+++ b/spec/lib.descope/api/v1/session_spec.rb
@@ -31,12 +31,32 @@
     end
 
     it 'is expected to post refresh session' do
-      jwt_response = { 'fake': 'response' }
-      allow(@instance).to receive(:generate_jwt_response).and_return(jwt_response)
-
-      expect(@instance).to receive(:post).with(REFRESH_TOKEN_PATH, {}, {}, 'refresh_token')
-      allow(@instance).to receive(:validate_token).with('refresh_token', nil).and_return({})
-      expect { @instance.refresh_session(refresh_token: 'refresh_token') }.not_to raise_error
+      jwt_response = {
+        'sessionJwt' => 'fake_session_jwt',
+        'refreshJwt' => 'fake_refresh_jwt',
+        'cookies' => {
+          'refresh_token' => 'fake_refresh_cookie'
+        }
+      }
+      refresh_token = 'refresh_token'
+      audience = nil
+
+      allow(@instance).to receive(:validate_refresh_token_not_nil).with(refresh_token).and_return(true)
+      allow(@instance).to receive(:validate_token).with(refresh_token, audience).and_return(true)
+      allow(@instance).to receive(:post).with(REFRESH_TOKEN_PATH, {}, {}, refresh_token).and_return(jwt_response)
+      refresh_cookie = jwt_response['cookies'][REFRESH_SESSION_COOKIE_NAME] || jwt_response['refreshJwt']
+
+      allow(@instance).to receive(:generate_jwt_response).with(
+        response_body: jwt_response,
+        refresh_cookie:,
+        audience:
+      ).and_return(jwt_response)
+
+      expect { @instance.refresh_session(refresh_token:, audience:) }.not_to raise_error
+
+      # Optionally verify the response if needed
+      result = @instance.refresh_session(refresh_token:, audience:)
+      expect(result).to eq(jwt_response)
     end
   end
 
diff --git a/spec/support/client_config.rb b/spec/support/client_config.rb
index 3e7dcdb..46eb8f6 100644
--- a/spec/support/client_config.rb
+++ b/spec/support/client_config.rb
@@ -5,7 +5,6 @@ module Configuration
   module_function
 
   def config
-    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
     raise 'DESCOPE_PROJECT_ID is not set' if ENV['DESCOPE_PROJECT_ID'].nil?
 
     {