Save some time initialising a context every spec

Add a context_aware option for specs that need to run
in their own context

ee/spec/workers/update_all_mirrors_worker_spec.rb

@@ -27,7 +27,7 @@ RSpec.describe UpdateAllMirrorsWorker do
       worker.perform
     end
 
-    it 'removes metadata except correlation_id from the application context before scheduling mirrors' do
+    it 'removes metadata except correlation_id from the application context before scheduling mirrors', :context_aware do
       inner_context = nil
       outer_context = nil
 

spec/lib/gitlab/metrics/background_transaction_spec.rb

@@ -30,7 +30,7 @@ RSpec.describe Gitlab::Metrics::BackgroundTransaction do
 
   describe '#labels' do
     it 'provides labels with endpoint_id and feature_category' do
-      Labkit::Context.with_context(feature_category: 'projects', caller_id: 'TestWorker') do
+      Gitlab::ApplicationContext.with_raw_context(feature_category: 'projects', caller_id: 'TestWorker') do
         expect(transaction.labels).to eq({ endpoint_id: 'TestWorker', feature_category: 'projects' })
       end
     end
@@ -41,7 +41,7 @@ RSpec.describe Gitlab::Metrics::BackgroundTransaction do
       value = 1
       expect(prometheus_metric).to receive(metric_method).with({ endpoint_id: 'TestWorker', feature_category: 'projects' }, value)
 
-      Labkit::Context.with_context(feature_category: 'projects', caller_id: 'TestWorker') do
+      Gitlab::ApplicationContext.with_raw_context(feature_category: 'projects', caller_id: 'TestWorker') do
         transaction.send(metric_method, :test_metric, value)
       end
     end

spec/spec_helper.rb

@@ -310,10 +310,20 @@ RSpec.configure do |config|
     RequestStore.clear!
   end
 
-  config.around do |example|
-    # Wrap each example in it's own context to make sure the contexts don't
-    # leak
-    Gitlab::ApplicationContext.with_raw_context { example.run }
+  if ENV['SKIP_RSPEC_CONTEXT_WRAPPING']
+    config.around(:example, :context_aware) do |example|
+      # Wrap each example in it's own context to make sure the contexts don't
+      # leak
+      Gitlab::ApplicationContext.with_raw_context { example.run }
+    end
+  else
+    config.around do |example|
+      if [:controller, :request, :feature].include?(example.metadata[:type]) || example.metadata[:context_aware]
+        Gitlab::ApplicationContext.with_raw_context { example.run }
+      else
+        example.run
+      end
+    end
   end
 
   config.around do |example|

spec/support/shared_examples/lib/api/ci/runner_shared_examples.rb

 # frozen_string_literal: true
 
 RSpec.shared_examples 'API::CI::Runner application context metadata' do |api_route|
-  it 'contains correct context metadata' do
+  it 'contains correct context metadata', :context_aware do
     # Avoids popping the context from the thread so we can
     # check its content after the request.
     allow(Labkit::Context).to receive(:pop)

spec/support/shared_examples/requests/api/logging_application_context_shared_examples.rb

 # frozen_string_literal: true
 
 RSpec.shared_examples 'storing arguments in the application context' do
-  around do |example|
-    Gitlab::ApplicationContext.with_base_context { example.run }
-  end
-
-  it 'places the expected params in the application context' do
+  it 'places the expected params in the application context', :context_aware do
     # Stub the clearing of the context so we can validate it later
-    # The `around` block above makes sure we do clean it up later
     allow(Labkit::Context).to receive(:pop)
 
     subject
 
-    Gitlab::ApplicationContext.with_base_context do |context|
-      expect(context.to_h)
-        .to include(log_hash(expected_params))
-    end
+    expect(Gitlab::ApplicationContext.current).to include(log_hash(expected_params))
   end
 
   def log_hash(hash)