module ActiveRecord::TestFixtures

def access_fixture(fs_name, *fixture_names)
  force_reload = fixture_names.pop if fixture_names.last == true || fixture_names.last == :reload
  return_single_record = fixture_names.size == 1
  fixture_names = @loaded_fixtures[fs_name].fixtures.keys if fixture_names.empty?
  @fixture_cache[fs_name] ||= {}
  instances = fixture_names.map do |f_name|
    f_name = f_name.to_s if f_name.is_a?(Symbol)
    @fixture_cache[fs_name].delete(f_name) if force_reload
    if @loaded_fixtures[fs_name][f_name]
      @fixture_cache[fs_name][f_name] ||= @loaded_fixtures[fs_name][f_name].find
    else
      raise StandardError, "No fixture named '#{f_name}' found for fixture set '#{fs_name}'"
    end
  end
  return_single_record ? instances.first : instances
end

def active_record_fixture(fixture_set_name, *fixture_names)
  if fs_name = fixture_sets[fixture_set_name.name]
    access_fixture(fs_name, *fixture_names)
  else
    raise StandardError, "No fixture set named '#{fixture_set_name.inspect}'"
  end
end

def after_teardown # :nodoc:
  super
ensure
  teardown_fixtures
end

def before_setup # :nodoc:
  setup_fixtures
  super
end

def fixture(fixture_set_name, *fixture_names)
  active_record_fixture(fixture_set_name, *fixture_names)
end

def instantiate_fixtures
  if pre_loaded_fixtures
    raise RuntimeError, "Load fixtures before instantiating them." if ActiveRecord::FixtureSet.all_loaded_fixtures.empty?
    ActiveRecord::FixtureSet.instantiate_all_loaded_fixtures(self, load_instances?)
  else
    raise RuntimeError, "Load fixtures before instantiating them." if @loaded_fixtures.nil?
    @loaded_fixtures.each_value do |fixture_set|
      ActiveRecord::FixtureSet.instantiate_fixtures(self, fixture_set, load_instances?)
    end
  end
end

def invalidate_already_loaded_fixtures
  @@already_loaded_fixtures.clear
end

def load_fixtures(config)
  ActiveRecord::FixtureSet.create_fixtures(fixture_paths, fixture_table_names, fixture_class_names, config).index_by(&:name)
end

def load_instances?
  use_instantiated_fixtures != :no_instances
end

def method_missing(method, ...)
  if fixture_sets.key?(method.name)
    active_record_fixture(method, ...)
  else
    super
  end
end

def respond_to_missing?(method, include_private = false)
  if include_private && fixture_sets.key?(method.name)
    true
  else
    super
  end
end

def run_in_transaction?
  use_transactional_tests &&
    !self.class.uses_transaction?(name)
end

def setup_asynchronous_queries_session
  @_async_queries_session = ActiveRecord::Base.asynchronous_queries_tracker.start_session
end

def setup_fixtures(config = ActiveRecord::Base)
  if pre_loaded_fixtures && !use_transactional_tests
    raise RuntimeError, "pre_loaded_fixtures requires use_transactional_tests"
  end
  @fixture_cache = {}
  @fixture_cache_key = [self.class.fixture_table_names.dup, self.class.fixture_paths.dup, self.class.fixture_class_names.dup]
  @fixture_connection_pools = []
  @@already_loaded_fixtures ||= {}
  @connection_subscriber = nil
  @saved_pool_configs = Hash.new { |hash, key| hash[key] = {} }
  if run_in_transaction?
    # Load fixtures once and begin transaction.
    @loaded_fixtures = @@already_loaded_fixtures[@fixture_cache_key]
    unless @loaded_fixtures
      @@already_loaded_fixtures.clear
      @loaded_fixtures = @@already_loaded_fixtures[@fixture_cache_key] = load_fixtures(config)
    end
    setup_transactional_fixtures
  else
    # Load fixtures for every test.
    ActiveRecord::FixtureSet.reset_cache
    invalidate_already_loaded_fixtures
    @loaded_fixtures = load_fixtures(config)
  end
  setup_asynchronous_queries_session
  # Instantiate fixtures for every test if requested.
  instantiate_fixtures if use_instantiated_fixtures
end

def setup_shared_connection_pool
  handler = ActiveRecord::Base.connection_handler
  handler.connection_pool_names.each do |name|
    pool_manager = handler.send(:connection_name_to_pool_manager)[name]
    pool_manager.shard_names.each do |shard_name|
      writing_pool_config = pool_manager.get_pool_config(ActiveRecord.writing_role, shard_name)
      @saved_pool_configs[name][shard_name] ||= {}
      pool_manager.role_names.each do |role|
        next unless pool_config = pool_manager.get_pool_config(role, shard_name)
        next if pool_config == writing_pool_config
        @saved_pool_configs[name][shard_name][role] = pool_config
        pool_manager.set_pool_config(role, shard_name, writing_pool_config)
      end
    end
  end
end

def setup_transactional_fixtures
  setup_shared_connection_pool
  # Begin transactions for connections already established
  @fixture_connection_pools = ActiveRecord::Base.connection_handler.connection_pool_list(:writing)
  @fixture_connection_pools.each do |pool|
    pool.pin_connection!(lock_threads)
    pool.lease_connection
  end
  # When connections are established in the future, begin a transaction too
  @connection_subscriber = ActiveSupport::Notifications.subscribe("!connection.active_record") do |_, _, _, _, payload|
    connection_name = payload[:connection_name] if payload.key?(:connection_name)
    shard = payload[:shard] if payload.key?(:shard)
    if connection_name
      pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(connection_name, shard: shard)
      if pool
        setup_shared_connection_pool
        unless @fixture_connection_pools.include?(pool)
          pool.pin_connection!(lock_threads)
          pool.lease_connection
          @fixture_connection_pools << pool
        end
      end
    end
  end
end

def teardown_asynchronous_queries_session
  ActiveRecord::Base.asynchronous_queries_tracker.finalize_session(true) if @_async_queries_session
end

def teardown_fixtures
  teardown_asynchronous_queries_session
  # Rollback changes if a transaction is active.
  if run_in_transaction?
    teardown_transactional_fixtures
  else
    ActiveRecord::FixtureSet.reset_cache
    invalidate_already_loaded_fixtures
  end
  ActiveRecord::Base.connection_handler.clear_active_connections!(:all)
end

def teardown_shared_connection_pool
  handler = ActiveRecord::Base.connection_handler
  @saved_pool_configs.each_pair do |name, shards|
    pool_manager = handler.send(:connection_name_to_pool_manager)[name]
    shards.each_pair do |shard_name, roles|
      roles.each_pair do |role, pool_config|
        next unless pool_manager.get_pool_config(role, shard_name)
        pool_manager.set_pool_config(role, shard_name, pool_config)
      end
    end
  end
  @saved_pool_configs.clear
end

def teardown_transactional_fixtures
  ActiveSupport::Notifications.unsubscribe(@connection_subscriber) if @connection_subscriber
  unless @fixture_connection_pools.map(&:unpin_connection!).all?
    # Something caused the transaction to be committed or rolled back
    # We can no longer trust the database is in a clean state.
    @@already_loaded_fixtures.clear
  end
  @fixture_connection_pools.clear
  teardown_shared_connection_pool
end