1 files changed, 20 insertions, 12 deletions

diff --git a/spec/initializers/diagnostic_reports_spec.rb b/spec/initializers/diagnostic_reports_spec.rb
index 01b1ed9b7b5..dc989efe809 100644
--- a/spec/initializers/diagnostic_reports_spec.rb
+++ b/spec/initializers/diagnostic_reports_spec.rb
@@ -2,15 +2,17 @@
 
 require 'spec_helper'
 
-RSpec.describe 'diagnostic reports' do
+RSpec.describe 'diagnostic reports', :aggregate_failures, feature_category: :application_performance do
   subject(:load_initializer) do
     load Rails.root.join('config/initializers/diagnostic_reports.rb')
   end
 
-  shared_examples 'does not modify worker startup hooks' do
+  shared_examples 'does not modify worker hooks' do
     it do
       expect(Gitlab::Cluster::LifecycleEvents).not_to receive(:on_worker_start)
+      expect(Gitlab::Cluster::LifecycleEvents).not_to receive(:on_worker_stop)
       expect(Gitlab::Memory::ReportsDaemon).not_to receive(:instance)
+      expect(Gitlab::Memory::Reporter).not_to receive(:new)
 
       load_initializer
     end
@@ -27,21 +29,27 @@ RSpec.describe 'diagnostic reports' do
       end
 
       let(:report_daemon) { instance_double(Gitlab::Memory::ReportsDaemon) }
+      let(:reporter) { instance_double(Gitlab::Memory::Reporter) }
 
       it 'modifies worker startup hooks, starts Gitlab::Memory::ReportsDaemon' do
         expect(Gitlab::Cluster::LifecycleEvents).to receive(:on_worker_start).and_call_original
-
+        expect(Gitlab::Cluster::LifecycleEvents).to receive(:on_worker_stop) # stub this out to not mutate global state
         expect_next_instance_of(Gitlab::Memory::ReportsDaemon) do |daemon|
-          expect(daemon).to receive(:start).and_call_original
+          expect(daemon).to receive(:start)
+        end
 
-          # make sleep no-op
-          allow(daemon).to receive(:sleep).and_return(nil)
+        load_initializer
+      end
 
-          # let alive return 3 times: true, true, false
-          allow(daemon).to receive(:alive).and_return(true, true, false)
-        end
+      it 'writes scheduled heap dumps in on_worker_stop' do
+        expect(Gitlab::Cluster::LifecycleEvents).to receive(:on_worker_start)
+        expect(Gitlab::Cluster::LifecycleEvents).to receive(:on_worker_stop).and_call_original
+        expect(Gitlab::Memory::Reporter).to receive(:new).and_return(reporter)
+        expect(reporter).to receive(:run_report).with(an_instance_of(Gitlab::Memory::Reports::HeapDump))
 
         load_initializer
+        # This is necessary because this hook normally fires during worker shutdown.
+        Gitlab::Cluster::LifecycleEvents.do_worker_stop
       end
     end
 
@@ -50,7 +58,7 @@ RSpec.describe 'diagnostic reports' do
         allow(::Gitlab::Runtime).to receive(:puma?).and_return(false)
       end
 
-      include_examples 'does not modify worker startup hooks'
+      include_examples 'does not modify worker hooks'
     end
   end
 
@@ -59,7 +67,7 @@ RSpec.describe 'diagnostic reports' do
       allow(::Gitlab::Runtime).to receive(:puma?).and_return(true)
     end
 
-    include_examples 'does not modify worker startup hooks'
+    include_examples 'does not modify worker hooks'
   end
 
   context 'when GITLAB_DIAGNOSTIC_REPORTS_ENABLED is set to false' do
@@ -68,6 +76,6 @@ RSpec.describe 'diagnostic reports' do
       allow(::Gitlab::Runtime).to receive(:puma?).and_return(true)
     end
 
-    include_examples 'does not modify worker startup hooks'
+    include_examples 'does not modify worker hooks'
   end
 end