Narrow the tracing of object allocations to during each test

1 files changed, 25 insertions, 25 deletions

diff --git a/tool/lib/minitest/unit.rb b/tool/lib/minitest/unit.rb
index ec6d6357b0..2cda29ad7d 100644
--- a/tool/lib/minitest/unit.rb
+++ b/tool/lib/minitest/unit.rb
@@ -961,37 +961,37 @@ module MiniTest
       }
 
       leakchecker = LeakChecker.new
+      if ENV["LEAK_CHECKER_TRACE_OBJECT_ALLOCATION"]
+        require "objspace"
+        trace = true
+      end
 
-      continuation = proc do
-        assertions = filtered_test_methods.map { |method|
-          inst = suite.new method
-          inst._assertions = 0
-
-          print "#{suite}##{method} = " if @verbose
-
-          start_time = Time.now if @verbose
-          result = inst.run self
+      assertions = filtered_test_methods.map { |method|
+        inst = suite.new method
+        inst._assertions = 0
 
-          print "%.2f s = " % (Time.now - start_time) if @verbose
-          print result
-          puts if @verbose
-          $stdout.flush
+        print "#{suite}##{method} = " if @verbose
 
-          unless defined?(RubyVM::JIT) && RubyVM::JIT.enabled? # compiler process is wrongly considered as leak
-            leakchecker.check("#{inst.class}\##{inst.__name__}")
+        start_time = Time.now if @verbose
+        result =
+          if trace
+            ObjectSpace.trace_object_allocations {inst.run self}
+          else
+            inst.run self
           end
 
-          inst._assertions
-        }
-        return assertions.size, assertions.inject(0) { |sum, n| sum + n }
-      end
+        print "%.2f s = " % (Time.now - start_time) if @verbose
+        print result
+        puts if @verbose
+        $stdout.flush
 
-      if ENV["LEAK_CHECKER_TRACE_OBJECT_ALLOCATION"]
-        require "objspace"
-        ObjectSpace.trace_object_allocations(&continuation)
-      else
-        continuation.call
-      end
+        unless defined?(RubyVM::JIT) && RubyVM::JIT.enabled? # compiler process is wrongly considered as leak
+          leakchecker.check("#{inst.class}\##{inst.__name__}")
+        end
+
+        inst._assertions
+      }
+      return assertions.size, assertions.inject(0) { |sum, n| sum + n }
     end
 
     ##