diff --git a/lib/mattermost/mattermost.rb b/lib/mattermost/session.rb
similarity index 95%
rename from lib/mattermost/mattermost.rb
rename to lib/mattermost/session.rb
index 84d016bb1979ec937f4c000edafc1d0d0f6ead52..d14121c91a018cb5f1addd82eed7af9eb9b507c3 100644
--- a/lib/mattermost/mattermost.rb
+++ b/lib/mattermost/session.rb
@@ -13,13 +13,14 @@ 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
 
     attr_accessor :current_resource_owner
 
     def initialize(uri, current_user)
+      # Sets the base uri for HTTParty, so we can use paths
       self.class.base_uri(uri)
 
       @current_resource_owner = current_user
@@ -27,8 +28,10 @@ module Mattermost
 
     def with_session
       raise NoSessionError unless create
-      yield
+      result = yield
       destroy
+
+      result
     end
 
     # Next methods are needed for Doorkeeper
@@ -85,7 +88,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
diff --git a/spec/lib/mattermost/mattermost_spec.rb b/spec/lib/mattermost/mattermost_spec.rb
deleted file mode 100644
index 7c99b4df9f3dbd6370d6833c8e80e8b960276f0b..0000000000000000000000000000000000000000
--- a/spec/lib/mattermost/mattermost_spec.rb
+++ /dev/null
@@ -1,42 +0,0 @@
-require 'spec_helper'
-
-describe Mattermost::Mattermost do
-  let(:user) { create(:user) }
-
-  subject { described_class.new('http://localhost:8065', user) }
-
-  # Needed for doorman to function
-  it { is_expected.to respond_to(:current_resource_owner) }
-  it { is_expected.to respond_to(:request) }
-  it { is_expected.to respond_to(:authorization) }
-  it { is_expected.to respond_to(:strategy) }
-
-  describe '#with session' do
-    let!(:stub) do
-      WebMock.stub_request(:get, 'http://localhost:8065/api/v3/oauth/gitlab/login').
-        to_return(headers: { 'location' => 'http://mylocation.com' }, status: 307)
-    end
-
-    context 'without oauth uri' do
-      it 'makes a request to the oauth uri' do
-        expect { subject.with_session }.to raise_error(Mattermost::NoSessionError)
-      end
-
-      context 'with oauth_uri' do
-        let!(:doorkeeper) do
-          Doorkeeper::Application.create(name: "GitLab Mattermost",
-                                         redirect_uri: "http://localhost:8065/signup/gitlab/complete\nhttp://localhost:8065/login/gitlab/complete",
-                                         scopes: "")
-        end
-
-        context 'without token_uri' do
-          it 'can not create a session' do
-            expect do
-              subject.with_session
-            end.to raise_error(Mattermost::NoSessionError)
-          end
-        end
-      end
-    end
-  end
-end
diff --git a/spec/lib/mattermost/session_spec.rb b/spec/lib/mattermost/session_spec.rb
new file mode 100644
index 0000000000000000000000000000000000000000..a93bab877dab7057838024a9b82b06b95f2342e0
--- /dev/null
+++ b/spec/lib/mattermost/session_spec.rb
@@ -0,0 +1,68 @@
+require 'spec_helper'
+
+describe Mattermost::Session do
+  let(:user) { create(:user) }
+
+  subject { described_class.new('http://localhost:8065', user) }
+
+  # Needed for doorkeeper to function
+  it { is_expected.to respond_to(:current_resource_owner) }
+  it { is_expected.to respond_to(:request) }
+  it { is_expected.to respond_to(:authorization) }
+  it { is_expected.to respond_to(:strategy) }
+
+  describe '#with session' do
+    let(:location) { 'http://location.tld' }
+    let!(:stub) do
+      WebMock.stub_request(:get, 'http://localhost:8065/api/v3/oauth/gitlab/login').
+        to_return(headers: { 'location' => location }, status: 307)
+    end
+
+    context 'without oauth uri' do
+      it 'makes a request to the oauth uri' do
+        expect { subject.with_session }.to raise_error(Mattermost::NoSessionError)
+      end
+    end
+
+    context 'with oauth_uri' do
+      let!(:doorkeeper) do
+        Doorkeeper::Application.create(name: "GitLab Mattermost",
+                                       redirect_uri: "http://localhost:8065/signup/gitlab/complete\nhttp://localhost:8065/login/gitlab/complete",
+                                       scopes: "")
+      end
+
+      context 'without token_uri' do
+        it 'can not create a session' do
+          expect do
+            subject.with_session
+          end.to raise_error(Mattermost::NoSessionError)
+        end
+      end
+
+      context 'with token_uri' do
+        let(:state) { "eyJhY3Rpb24iOiJsb2dpbiIsImhhc2giOiIkMmEkMTAkVC9wYVlEaTdIUS8vcWdKRmdOOUllZUptaUNJWUlvNVNtNEcwU2NBMXFqelNOVmVPZ1cxWUsifQ%3D%3D" }
+        let(:location) { "http://locahost:8065/oauth/authorize?response_type=code&client_id=#{doorkeeper.uid}&redirect_uri=http%3A%2F%2Flocalhost:8065%2Fsignup%2Fgitlab%2Fcomplete&state=#{state}" }
+
+        before do
+          WebMock.stub_request(:get, /http:\/\/localhost:8065\/signup\/gitlab\/complete*/).
+            to_return(headers: { 'token' => 'thisworksnow' }, status: 202)
+        end
+
+        it 'can setup a session' do
+          expect(subject).to receive(:destroy)
+
+          subject.with_session { 1 + 1 }
+        end
+
+        it 'returns the value of the block' do
+          WebMock.stub_request(:post, "http://localhost:8065/api/v3/users/logout").
+            to_return(headers: { 'token' => 'thisworksnow' }, status: 200)
+
+          value = subject.with_session { 1 + 1 }
+
+          expect(value).to be(2)
+        end
+      end
+    end
+  end
+end