test_unit_profiler.rb

require 'test/unit/ui/console/testrunner'

class TestUnitProfiler
  Record = Struct.new(:test_name, :test_class, :start_time, :self_time, :total_time, :children, keyword_init: true) do
    def to_h
      {
        name: test_name,
        class: test_class,
        self: self_time,
        total: total_time,
        children: children.map(&:to_h),
      }
    end

    def to_json
      to_h.to_json
    end
  end

  def initialize
    @root = Record.new(children: [])
    @stack = [@root]
  end

  def result
    @root.children.first.children.map(&:to_h)
  end

  def enter(test)
    @stack << Record.new(
      test_name: test.name,
      test_class: test.class.name,
      start_time: Process.clock_gettime(Process::CLOCK_MONOTONIC),
      children: [],
    )
  end

  def exit
    end_time = Process.clock_gettime(Process::CLOCK_MONOTONIC)
    profile = @stack.pop
    profile.total_time = end_time - profile.start_time
    profile.self_time = profile.total_time - profile.children.sum(&:total_time)
    @stack.last.children << profile
  end
end

module TestUnitProfile
  def attach_to_mediator
    super
    @profiler = TestUnitProfiler.new
    @mediator.add_listener(Test::Unit::TestSuite::STARTED_OBJECT, &method(:test_unit_profile_enter))
    @mediator.add_listener(Test::Unit::TestSuite::FINISHED_OBJECT, &method(:test_unit_profile_exit))
    @mediator.add_listener(Test::Unit::TestCase::STARTED_OBJECT, &method(:test_unit_profile_enter))
    @mediator.add_listener(Test::Unit::TestCase::FINISHED_OBJECT, &method(:test_unit_profile_exit))
    @mediator.add_listener(Test::Unit::UI::TestRunnerMediator::FINISHED, &method(:test_unit_profile_output_result))
  end

  def test_unit_profile_enter(test)
    @profiler.enter(test)
  end

  def test_unit_profile_exit(test)
    @profiler.exit
  end

  def test_unit_profile_output_result(elapsed_time)
    result_json = @profiler.result.to_json
    if ENV["TEST_UNIT_PROFILE_OUTPUT"]
      File.write(ENV["TEST_UNIT_PROFILE_OUTPUT"], result_json)
    else
      puts result_json
    end
  end
end

class Test::Unit::UI::Console::TestRunner
  prepend TestUnitProfile
end