class RuboCop::Cop::Style::ObjectThen

def check_method_node(node)
  return unless preferred_method(node)
  message = message(node)
  add_offense(node.loc.selector, message: message) do |corrector|
    corrector.replace(node.loc.selector, style.to_s)
  end
end

def message(node)
  format(MSG, prefer: style.to_s, current: node.method_name)
end

def on_block(node)
  check_method_node(node.send_node)
end

def on_send(node)
  return unless node.arguments.one? && node.first_argument.block_pass_type?
  check_method_node(node)
end

def preferred_method(node)
  case style
  when :then
    node.method?(:yield_self)
  when :yield_self
    node.method?(:then)
  else
    false
  end
end