diff --git a/.travis.yml b/.travis.yml
index b562811..e795a25 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -16,4 +16,4 @@ gemfile:
   - Gemfile
   - gemfiles/mongoid2.gemfile
   - gemfiles/mongoid3.gemfile
-  - gemfiles/mongoid4.gemfile
\ No newline at end of file
+  - gemfiles/mongoid4.gemfile
diff --git a/lib/searchkick/results.rb b/lib/searchkick/results.rb
index 47e7023..2c469e6 100644
--- a/lib/searchkick/results.rb
+++ b/lib/searchkick/results.rb
@@ -27,12 +27,14 @@ module Searchkick
               records = records.includes(options[:includes])
             end
             results[type] =
-              if records.respond_to?(:primary_key) && records.primary_key
+              if records.respond_to?(:primary_key) and records.primary_key
+                # ActiveRecord
                 records.where(records.primary_key => grouped_hits.map{|hit| hit["_id"] }).to_a
-              elsif records.respond_to?(:all) && records.all && records.all.respond_to?(:for_ids)
+              elsif records.respond_to?(:all) and records.all.respond_to?(:for_ids)
                 # Mongoid 2
                 records.all.for_ids(grouped_hits.map{|hit| hit["_id"] }).to_a
               else
+                # Mongoid 3+
                 records.queryable.for_ids(grouped_hits.map{|hit| hit["_id"] }).to_a
               end
           end
diff --git a/test/test_helper.rb b/test/test_helper.rb
index 39eb54a..edf974e 100644
--- a/test/test_helper.rb
+++ b/test/test_helper.rb
@@ -11,15 +11,12 @@ Searchkick.client.transport.logger = Logger.new("elasticsearch.log")
 
 if defined?(Mongoid)
 
-  # Helpers to determine version of mongoid.
-  module Mongoid
-    def self.mongoid2?
-      ::Mongoid.const_defined? :Contexts # deprecated in Mongoid 3.x
-    end
+  def mongoid2?
+    Mongoid::VERSION.starts_with?("2.")
   end
 
-  # Mongoid2 uses BSON gem. We need to re-define <=> in order for TestSql.test_order_id to pass.
-  if Mongoid.mongoid2? && defined?(BSON) && defined?(BSON::ObjectId)
+  if mongoid2?
+    # enable comparison of BSON::ObjectIds
     module BSON
       class ObjectId
         def <=>(other)
@@ -30,7 +27,7 @@ if defined?(Mongoid)
   end
 
   Mongoid.configure do |config|
-    if Mongoid.mongoid2?
+    if mongoid2?
       config.master = Mongo::Connection.new.db("searchkick_test")
     else
       config.connect_to "searchkick_test"
--
libgit2 0.21.0