1 files changed, 15 insertions, 1 deletions
diff --git a/spec/lib/gitlab/application_context_spec.rb b/spec/lib/gitlab/application_context_spec.rb
index 0fbbc67ef6a..c4fe2ebaba9 100644
--- a/spec/lib/gitlab/application_context_spec.rb
+++ b/spec/lib/gitlab/application_context_spec.rb
@@ -27,6 +27,20 @@ RSpec.describe Gitlab::ApplicationContext do
     end
   end
 
+  describe '.with_raw_context' do
+    it 'yields the block' do
+      expect { |b| described_class.with_raw_context({}, &b) }.to yield_control
+    end
+
+    it 'passes the attributes unaltered on to labkit' do
+      attrs = { foo: :bar }
+
+      expect(Labkit::Context).to receive(:with_context).with(attrs)
+
+      described_class.with_raw_context(attrs) {}
+    end
+  end
+
   describe '.push' do
     it 'passes the expected context on to labkit' do
       fake_proc = duck_type(:call)
@@ -138,7 +152,7 @@ RSpec.describe Gitlab::ApplicationContext do
     it 'does not cause queries' do
       context = described_class.new(project: create(:project), namespace: create(:group, :nested), user: create(:user))
 
-      expect { context.use { Labkit::Context.current.to_h } }.not_to exceed_query_limit(0)
+      expect { context.use { Gitlab::ApplicationContext.current } }.not_to exceed_query_limit(0)
     end
   end
 end