class ActiveSupport::Testing::SimpleStubs

def initialize
  @stubs = Concurrent::Map.new { |h, k| h[k] = {} }
end

def stub_object(object, method_name, &block)
  if stub = stubbing(object, method_name)
    unstub_object(stub)
  end
  new_name = "__simple_stub__#{method_name}"
  @stubs[object.object_id][method_name] = Stub.new(object, method_name, new_name)
  object.singleton_class.alias_method new_name, method_name
  object.define_singleton_method(method_name, &block)
end

def stubbing(object, method_name)
  @stubs[object.object_id][method_name]
end

def unstub_all!
  @stubs.each_value do |object_stubs|
    object_stubs.each_value do |stub|
      unstub_object(stub)
    end
  end
  @stubs.clear
end

def unstub_object(stub)
  singleton_class = stub.object.singleton_class
  singleton_class.silence_redefinition_of_method stub.method_name
  singleton_class.alias_method stub.method_name, stub.original_method
  singleton_class.undef_method stub.original_method
end