Improve session tests

1 files changed, 5 insertions, 3 deletions

diff --git a/lib/mattermost/mattermost.rb b/lib/mattermost/session.rb
index 84d016bb197..81964666757 100644
--- a/lib/mattermost/mattermost.rb
+++ b/lib/mattermost/session.rb
@@ -13,7 +13,7 @@ module Mattermost
   # This class however skips the button click, and also the approval phase to
   # speed up the process and keep it without manual action and get a session
   # going.
-  class Mattermost
+  class Session
     include Doorkeeper::Helpers::Controller
     include HTTParty
 
@@ -27,8 +27,10 @@ module Mattermost
 
     def with_session
       raise NoSessionError unless create
-      yield
+      result = yield
       destroy
+
+      result
     end
 
     # Next methods are needed for Doorkeeper
@@ -85,7 +87,7 @@ module Mattermost
     end
 
     def request_token
-      @request_token ||= if @token_uri
+      @request_token ||= begin
                            response = get(@token_uri, follow_redirects: false)
                            response.headers['token'] if 200 <= response.code && response.code < 400
                          end