1 files changed, 8 insertions, 12 deletions
diff --git a/spec/lib/mattermost/session_spec.rb b/spec/lib/mattermost/session_spec.rb
index 5410bfbeb31..b7687d48c68 100644
--- a/spec/lib/mattermost/session_spec.rb
+++ b/spec/lib/mattermost/session_spec.rb
@@ -1,6 +1,8 @@
 require 'spec_helper'
 
 describe Mattermost::Session, type: :request do
+  include ExclusiveLeaseHelpers
+
   let(:user) { create(:user) }
 
   let(:gitlab_url) { "http://gitlab.com" }
@@ -97,26 +99,20 @@ describe Mattermost::Session, type: :request do
       end
     end
 
-    context 'with lease' do
-      before do
-        allow(subject).to receive(:lease_try_obtain).and_return('aldkfjsldfk')
-      end
+    context 'exclusive lease' do
+      let(:lease_key) { 'mattermost:session' }
 
       it 'tries to obtain a lease' do
-        expect(subject).to receive(:lease_try_obtain)
-        expect(Gitlab::ExclusiveLease).to receive(:cancel)
+        expect_to_obtain_exclusive_lease(lease_key, 'uuid')
+        expect_to_cancel_exclusive_lease(lease_key, 'uuid')
 
         # Cannot setup a session, but we should still cancel the lease
         expect { subject.with_session }.to raise_error(Mattermost::NoSessionError)
       end
-    end
 
-    context 'without lease' do
-      before do
-        allow(subject).to receive(:lease_try_obtain).and_return(nil)
-      end
+      it 'returns a NoSessionError error without lease' do
+        stub_exclusive_lease_taken(lease_key)
 
-      it 'returns a NoSessionError error' do
         expect { subject.with_session }.to raise_error(Mattermost::NoSessionError)
       end
     end