lib/elasticity/multi_search.rb in es-elasticity-0.4.5 vs lib/elasticity/multi_search.rb in es-elasticity-0.5.0

@@ -1,7 +1,8 @@
 module Elasticity
   class MultiSearch
+
     def initialize
       @searches = {}
       @mappers  = {}
       yield self if block_given?
     end
@@ -21,35 +22,37 @@
 
       name
     end
 
     def [](name)
-      @results ||= fetch
-      @results[name]
+      results_collection[name]
     end
 
     private
 
+    def results_collection
+      @results_collection ||= fetch
+    end
+
     def fetch
       bodies = @searches.values.map do |hsh|
         hsh[:search_definition].to_msearch_args
       end
 
       response = ActiveSupport::Notifications.instrument("multi_search.elasticity", args: { body: bodies }) do
         Elasticity.config.client.msearch(body: bodies.map(&:dup))
       end
-
       results = {}
 
       @searches.keys.each_with_index do |name, idx|
         resp = response["responses"][idx]
         search = @searches[name]
 
         results[name] = case
         when search[:documents]
-          resp["hits"]["hits"].map { |hit| search[:documents].map_hit(hit) }
+          Search::Results.new(resp, search[:search_definition].body, search[:documents].method(:map_hit))
         when search[:active_records]
-          Search::ActiveRecordProxy.from_hits(search[:active_records], resp["hits"]["hits"])
+          Search::ActiveRecordProxy.map_response(search[:active_records], search[:search_definition].body, resp)
         end
       end
 
       results
     end