module ActiveRecord::TestFixtures

def after_teardown # :nodoc:
  super
  teardown_fixtures
end

def before_setup # :nodoc:
  setup_fixtures
  super
end

def enlist_fixture_connections
  setup_shared_connection_pool
  ActiveRecord::Base.connection_handler.connection_pool_list.map(&:connection)
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 load_fixtures(config)
  ActiveRecord::FixtureSet.create_fixtures(fixture_path, fixture_table_names, fixture_class_names, config).index_by(&:name)
end

def load_instances?
  use_instantiated_fixtures != :no_instances
end

def run_in_transaction?
  use_transactional_tests &&
    !self.class.uses_transaction?(name)
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_connections = []
  @@already_loaded_fixtures ||= {}
  @connection_subscriber = nil
  @legacy_saved_pool_configs = Hash.new { |hash, key| hash[key] = {} }
  @saved_pool_configs = Hash.new { |hash, key| hash[key] = {} }
  # Load fixtures once and begin transaction.
  if run_in_transaction?
    if @@already_loaded_fixtures[self.class]
      @loaded_fixtures = @@already_loaded_fixtures[self.class]
    else
      @loaded_fixtures = load_fixtures(config)
      @@already_loaded_fixtures[self.class] = @loaded_fixtures
    end
    # Begin transactions for connections already established
    @fixture_connections = enlist_fixture_connections
    @fixture_connections.each do |connection|
      connection.begin_transaction joinable: false, _lazy: false
      connection.pool.lock_thread = true if lock_threads
    end
    # When connections are established in the future, begin a transaction too
    @connection_subscriber = ActiveSupport::Notifications.subscribe("!connection.active_record") do |_, _, _, _, payload|
      spec_name = payload[:spec_name] if payload.key?(:spec_name)
      shard = payload[:shard] if payload.key?(:shard)
      setup_shared_connection_pool if ActiveRecord.legacy_connection_handling
      if spec_name
        begin
          connection = ActiveRecord::Base.connection_handler.retrieve_connection(spec_name, shard: shard)
        rescue ConnectionNotEstablished
          connection = nil
        end
        if connection
          setup_shared_connection_pool unless ActiveRecord.legacy_connection_handling
          if !@fixture_connections.include?(connection)
            connection.begin_transaction joinable: false, _lazy: false
            connection.pool.lock_thread = true if lock_threads
            @fixture_connections << connection
          end
        end
      end
    end
  # Load fixtures for every test.
  else
    ActiveRecord::FixtureSet.reset_cache
    @@already_loaded_fixtures[self.class] = nil
    @loaded_fixtures = load_fixtures(config)
  end
  # Instantiate fixtures for every test if requested.
  instantiate_fixtures if use_instantiated_fixtures
end

def setup_shared_connection_pool
  if ActiveRecord.legacy_connection_handling
    writing_handler = ActiveRecord::Base.connection_handlers[ActiveRecord.writing_role]
    ActiveRecord::Base.connection_handlers.values.each do |handler|
      if handler != writing_handler
        handler.connection_pool_names.each do |name|
          writing_pool_manager = writing_handler.send(:owner_to_pool_manager)[name]
          return unless writing_pool_manager
          pool_manager = handler.send(:owner_to_pool_manager)[name]
          @legacy_saved_pool_configs[handler][name] ||= {}
          pool_manager.shard_names.each do |shard_name|
            writing_pool_config = writing_pool_manager.get_pool_config(nil, shard_name)
            pool_config = pool_manager.get_pool_config(nil, shard_name)
            next if pool_config == writing_pool_config
            @legacy_saved_pool_configs[handler][name][shard_name] = pool_config
            pool_manager.set_pool_config(nil, shard_name, writing_pool_config)
          end
        end
      end
    end
  else
    handler = ActiveRecord::Base.connection_handler
    handler.connection_pool_names.each do |name|
      pool_manager = handler.send(:owner_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
end

def teardown_fixtures
  # Rollback changes if a transaction is active.
  if run_in_transaction?
    ActiveSupport::Notifications.unsubscribe(@connection_subscriber) if @connection_subscriber
    @fixture_connections.each do |connection|
      connection.rollback_transaction if connection.transaction_open?
      connection.pool.lock_thread = false
    end
    @fixture_connections.clear
    teardown_shared_connection_pool
  else
    ActiveRecord::FixtureSet.reset_cache
  end
  ActiveRecord::Base.clear_active_connections!
end

def teardown_shared_connection_pool
  if ActiveRecord.legacy_connection_handling
    @legacy_saved_pool_configs.each_pair do |handler, names|
      names.each_pair do |name, shards|
        shards.each_pair do |shard_name, pool_config|
          pool_manager = handler.send(:owner_to_pool_manager)[name]
          pool_manager.set_pool_config(nil, shard_name, pool_config)
        end
      end
    end
  else
    handler = ActiveRecord::Base.connection_handler
    @saved_pool_configs.each_pair do |name, shards|
      pool_manager = handler.send(:owner_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
  end
  @legacy_saved_pool_configs.clear
  @saved_pool_configs.clear
end