Test::Unitのtestやsetup/cleanup/teardownのプロファイルをとる(Rails編)

test_unit_profiling_for_rails.rb

# test/test_helper.rbに書く

module TestProfiling
  module TestCase
    def run(result)
      data = {
        name: method_name,
        owner_name: self.class.name,
        source_location: method(method_name).source_location,
        passed: passed?,
      }
      ActiveSupport::Notifications.instrument("test_unit.run_test_case", data) do
        with_benchmark(data) { super }
      end
    end

    private

    def run_fixture_callback(method_name, options, &block)
      method = method(method_name)
      data = {
        name: method_name,
        owner_name: method.owner.name,
        source_location: method.source_location,
      }
      ActiveSupport::Notifications.instrument("test_unit.run_fixture_callback", data) do
        with_benchmark(data) { super }
      end
    end

    def run_test
      data = {
        name: @method_name,
        owner_name: self.class.name,
        source_location: method(@method_name).source_location,
      }
      ActiveSupport::Notifications.instrument("test_unit.run_test", data) do
        with_benchmark(data) { super }
      end
    end

    def with_benchmark(data)
      ret = nil
      data[:elapsed_time] = Benchmark.realtime do
        ret = yield
      end
      ret
    end
  end
  ::Test::Unit::TestCase.prepend TestCase

  module TestSuite
    def run(result, &progress_block)
      data = {
        test_suite: name,
        passed: passed?,
      }
      ActiveSupport::Notifications.instrument("test_unit.run_test_suite", data) do
        super
      end
    end
  end
  ::Test::Unit::TestSuite.prepend TestSuite
end

max_result_count = 15
test_unit_data = {
  fixture: {},
  test: {},
}
test_root_regexp = %r{\A#{Regexp.quote(Rails.root.to_s)}/test/}

ActiveSupport::Notifications.subscribe(/\Atest_unit.run_(?:fixture_callback|test)\z/) do |name, _started, _finished, _unique_id, data|
  source_location = data[:source_location]
  source_file, = source_location
  next unless test_root_regexp.match?(source_file)
  next unless data[:elapsed_time]

  if name == "test_unit.run_fixture_callback"
    type = :fixture
    key = [data[:owner_name], *source_location]
  else
    type = :test
    key = [data[:owner_name], data[:name]]
  end

  test_unit_data[type][key] ||= {
    count: 0,
    elapsed_time: 0.0,
    name: data[:name],
    owner_name: data[:owner_name],
    source_location: source_location,
  }
  test_unit_data[type][key][:count] += 1
  test_unit_data[type][key][:elapsed_time] += data[:elapsed_time]
end

Test::Unit.at_exit do
  fixture_callback_name = lambda do |source_file, lineno, name|
    if /_\d+\z/ =~ name
      "#{Regexp.last_match.pre_match}(#{source_file.sub(test_root_regexp, '')}:#{lineno})"
    else
      name
    end
  end

  results = {
    test_by_owner: test_unit_data[:test].each_with_object({}) do |((owner_name, name), data), h|
      name = name.sub(/\Atest: /, "")
      h["#{owner_name}\##{name}"] = data
    end,
    fixture_by_owner: test_unit_data[:fixture].each_with_object({}) do |((owner_name, source_file, lineno), data), h|
      name_with_location = fixture_callback_name.call(source_file, lineno, data[:name])
      h["#{owner_name}\##{name_with_location}"] = data
    end,
    fixture_by_source: test_unit_data[:fixture].each_with_object({}) do |((_, source_file, lineno), data), h|
      name_with_location = fixture_callback_name.call(source_file, lineno, data[:name])
      if h[name_with_location]
        h[name_with_location][:count] += data[:count]
        h[name_with_location][:elapsed_time] += data[:elapsed_time]
      else
        h[name_with_location] = data.dup
        h[name_with_location].delete(:owner_name)
      end
    end,
  }

  results.each do |type, result|
    print "#{type}::\n\n"

    total_time = result.sum {|_, data| data[:elapsed_time] }
    result.sort_by {|_, data| -data[:elapsed_time] }.each_with_index do |(key, data), idx|
      break if idx == max_result_count

      printf(
        "%6.2fs %5.1f%% (%4d) %.2fs %s\n",
        data[:elapsed_time],
        data[:elapsed_time] * 100 / total_time,
        data[:count],
        data[:elapsed_time] / data[:count],
        key,
      )
    end

    print "\n\n"
  end
end