require 'mongoid/factory'

module RocketJob
  # Don't convert to Mongoid::Factory since it conflicts with Mongoid use.
  module MongoidFactory
    def from_db(klass, attributes = nil, criteria = nil)
      obj                 = super(klass, attributes, criteria)
      obj.collection_name = criteria.collection_name if criteria
      obj
    end
  end
end

::Mongoid::Factory.extend(RocketJob::MongoidFactory)