From e67fe2522f3c6b61d2f298abe04e8f52a541c8df Mon Sep 17 00:00:00 2001
From: ellnix <103502144+ellnix@users.noreply.github.com>
Date: Mon, 12 Feb 2024 16:30:39 +0100
Subject: [PATCH] Add MultiSearchResult

---
 lib/meilisearch/rails/multi_search.rb        |  51 +-------
 lib/meilisearch/rails/multi_search/result.rb |  86 ++++++++++++++
 lib/meilisearch/rails/utilities.rb           |   8 ++
 spec/multi_search/result_spec.rb             | 115 ++++++++++++++++++-
 4 files changed, 210 insertions(+), 50 deletions(-)
 create mode 100644 lib/meilisearch/rails/multi_search/result.rb

diff --git a/lib/meilisearch/rails/multi_search.rb b/lib/meilisearch/rails/multi_search.rb
index ed76427f..64d2abc4 100644
--- a/lib/meilisearch/rails/multi_search.rb
+++ b/lib/meilisearch/rails/multi_search.rb
@@ -1,3 +1,5 @@
+require_relative 'multi_search/result'
+
 module MeiliSearch
   module Rails
     class << self
@@ -15,54 +17,7 @@ def multi_search(searches)
 
         raw_results = client.multi_search(search_parameters)['results']
 
-        searches.zip(raw_results).flat_map do |(index_target, search_options), result|
-          index_target = search_options[:class_name].constantize if search_options[:class_name]
-
-          case index_target
-          when String, Symbol
-            result['hits']
-          else
-            load_results(index_target, result)
-          end
-        end
-      end
-
-      private
-
-      def load_results(klass, result)
-        pk_method = if defined?(::Mongoid::Document) && klass.include?(::Mongoid::Document)
-                      klass.ms_primary_key_method.in
-                    else
-                      klass.ms_primary_key_method
-                    end
-
-        ms_pk = klass.meilisearch_options[:primary_key] || IndexSettings::DEFAULT_PRIMARY_KEY
-
-        db_is_sequel = defined?(::Sequel::Model) && klass < Sequel::Model
-        pk_is_virtual = klass.columns.map(&(db_is_sequel ? :to_s : :name)).exclude?(pk_method.to_s)
-
-        condition_key = pk_is_virtual ? klass.primary_key : pk_method
-
-        hits_by_id =
-          result['hits'].index_by { |hit| hit[pk_is_virtual ? condition_key : ms_pk.to_s] }
-
-        records = klass.where(condition_key => hits_by_id.keys)
-
-        if records.respond_to? :in_order_of
-          records.in_order_of(pk_method, hits_by_id.keys).each do |record|
-            record.formatted = hits_by_id[record.send(pk_method).to_s]['_formatted']
-          end
-        else
-          results_by_id = records.index_by do |hit|
-            hit.send(pk_method).to_s
-          end
-
-          result['hits'].filter_map do |hit|
-            record = results_by_id[hit[ms_pk.to_s].to_s]
-            record&.formatted = hit['_formatted']
-            record
-          end
-        end
+        MultiSearchResult.new(searches, raw_results)
       end
     end
   end
diff --git a/lib/meilisearch/rails/multi_search/result.rb b/lib/meilisearch/rails/multi_search/result.rb
new file mode 100644
index 00000000..20118f7e
--- /dev/null
+++ b/lib/meilisearch/rails/multi_search/result.rb
@@ -0,0 +1,86 @@
+module MeiliSearch
+  module Rails
+    class MultiSearchResult
+      attr_reader :metadata
+
+      def initialize(searches, raw_results)
+        @results = {}
+        @metadata = {}
+
+        searches.zip(raw_results).each do |(index_target, search_options), result|
+          index_target = search_options[:class_name].constantize if search_options[:class_name]
+
+          @results[index_target] = case index_target
+                                   when String, Symbol
+                                     result['hits']
+                                   else
+                                     load_results(index_target, result)
+                                   end
+
+          @metadata[index_target] = result.except('hits')
+        end
+      end
+
+      include Enumerable
+
+      def each_hit
+        @results.each do |_index_target, results|
+          results.each { |res| yield res }
+        end
+      end
+      alias_method :each, :each_hit
+
+      def each_result
+        @results.each
+      end
+
+      def to_a
+        @results.values.flatten(1)
+      end
+      alias_method :to_ary, :to_a
+
+      def to_h
+        @results
+      end
+      alias_method :to_hash, :to_h
+
+      private
+
+      def load_results(klass, result)
+        pk_method = klass.ms_primary_key_method
+        pk_method = pk_method.in if Utilities.is_mongo_model?(klass)
+
+        ms_pk = klass.meilisearch_options[:primary_key] || IndexSettings::DEFAULT_PRIMARY_KEY
+
+        condition_key = pk_is_virtual?(klass, pk_method) ? klass.primary_key : pk_method
+
+        hits_by_id =
+          result['hits'].index_by { |hit| hit[condition_key.to_s] }
+
+        records = klass.where(condition_key => hits_by_id.keys)
+
+        if records.respond_to? :in_order_of
+          records.in_order_of(condition_key, hits_by_id.keys).each do |record|
+            record.formatted = hits_by_id[record.send(condition_key).to_s]['_formatted']
+          end
+        else
+          results_by_id = records.index_by do |hit|
+            hit.send(condition_key).to_s
+          end
+
+          result['hits'].filter_map do |hit|
+            record = results_by_id[hit[condition_key.to_s].to_s]
+            record&.formatted = hit['_formatted']
+            record
+          end
+        end
+      end
+
+      def pk_is_virtual?(model_class, pk_method)
+        model_class.columns
+          .map(&(Utilities.is_sequel_model?(model_class) ? :to_s : :name))
+          .exclude?(pk_method.to_s)
+      end
+    end
+  end
+end
diff --git a/lib/meilisearch/rails/utilities.rb b/lib/meilisearch/rails/utilities.rb
index 377eca3d..e15bdc91 100644
--- a/lib/meilisearch/rails/utilities.rb
+++ b/lib/meilisearch/rails/utilities.rb
@@ -48,6 +48,14 @@ def indexable?(record, options)
           true
         end
 
+        def is_mongo_model?(model_class)
+          defined?(::Mongoid::Document) && model_class.include?(::Mongoid::Document)
+        end
+
+        def is_sequel_model?(model_class)
+          defined?(::Sequel::Model) && model_class < Sequel::Model
+        end
+
         private
 
         def constraint_passes?(record, constraint)
diff --git a/spec/multi_search/result_spec.rb b/spec/multi_search/result_spec.rb
index ad3430dc..74372cf3 100644
--- a/spec/multi_search/result_spec.rb
+++ b/spec/multi_search/result_spec.rb
@@ -1,5 +1,116 @@
 require 'spec_helper'
 
-describe 'MeiliSearch::Rails::MultiSearchResult' do # rubocop:todo RSpec/EmptyExampleGroup
-  # TODO: Write specs
+describe MeiliSearch::Rails::MultiSearchResult do
+  it 'is enumerable' do
+    expect(described_class).to include(Enumerable)
+  end
+
+  let(:raw_results) do
+    [
+      { 'indexUid' => 'books_index',
+        'hits' => [{ 'name' => 'Steve Jobs', 'id' => '3', 'author' => 'Walter Isaacson', 'premium' => nil, 'released' => nil, 'genre' => nil }],
+        'query' => 'Steve', 'processingTimeMs' => 0, 'limit' => 20, 'offset' => 0, 'estimatedTotalHits' => 1
+      },
+      { 'indexUid' => 'products_index',
+        'hits' => [{ 'id' => '4', 'href' => 'ebay', 'name' => 'palm pixi plus' }],
+        'query' => 'palm', 'processingTimeMs' => 0, 'limit' => 1, 'offset' => 0, 'estimatedTotalHits' => 2
+      },
+      { 'indexUid' => 'color_index',
+        'hits' => [
+          { 'name' => 'black', 'id' => '5', 'short_name' => 'bla', 'hex' => 0 },
+          { 'name' => 'blue', 'id' => '4', 'short_name' => 'blu', 'hex' => 255 }
+        ],
+        'query' => 'bl', 'processingTimeMs' => 0, 'limit' => 20, 'offset' => 0, 'estimatedTotalHits' => 2
+      }
+    ]
+  end
+
+  context 'with index name keys' do
+    subject(:result) { described_class.new(searches, raw_results) }
+
+    let(:searches) do
+      {
+        'books_index' => { q: 'Steve' },
+        'products_index' => { q: 'palm', limit: 1 },
+        'color_index' => { q: 'bl' }
+      }
+    end
+
+    it 'enumerates through the hits' do
+      expect(result).to contain_exactly(
+        a_hash_including('author' => 'Walter Isaacson', 'name' => 'Steve Jobs'),
+        a_hash_including('name' => 'palm pixi plus'),
+        a_hash_including('name' => 'blue', 'short_name' => 'blu'),
+        a_hash_including('name' => 'black', 'short_name' => 'bla')
+      )
+    end
+
+    it 'enumerates through the hits of each result with #each_result' do
+      expect(result.each_result).to be_an(Enumerator)
+      expect(result.each_result).to contain_exactly(
+          [ 'books_index', contain_exactly(
+            a_hash_including('author' => 'Walter Isaacson', 'name' => 'Steve Jobs')) ],
+          [ 'products_index', contain_exactly(
+            a_hash_including('name' => 'palm pixi plus')) ],
+          [ 'color_index', contain_exactly(
+            a_hash_including('name' => 'blue', 'short_name' => 'blu'),
+            a_hash_including('name' => 'black', 'short_name' => 'bla')) ]
+      )
+    end
+
+    describe '#to_a' do
+      it 'returns the hits' do
+        expect(result.to_a).to contain_exactly(
+          a_hash_including('author' => 'Walter Isaacson', 'name' => 'Steve Jobs'),
+          a_hash_including('name' => 'palm pixi plus'),
+          a_hash_including('name' => 'blue', 'short_name' => 'blu'),
+          a_hash_including('name' => 'black', 'short_name' => 'bla')
+        )
+      end
+
+      it 'aliases as #to_ary' do
+        expect(subject.method(:to_ary).original_name).to eq :to_a
+      end
+    end
+
+    describe '#to_h' do
+      it 'returns a hash of indexes and hits' do
+        expect(result.to_h).to match(
+          'books_index' => contain_exactly(
+            a_hash_including('author' => 'Walter Isaacson', 'name' => 'Steve Jobs')
+          ),
+          'products_index' => contain_exactly(
+            a_hash_including('name' => 'palm pixi plus')
+          ),
+          'color_index' => contain_exactly(
+            a_hash_including('name' => 'blue', 'short_name' => 'blu'),
+            a_hash_including('name' => 'black', 'short_name' => 'bla')
+          )
+        )
+      end
+
+      it 'is aliased as #to_hash' do
+        expect(result.method(:to_hash).original_name).to eq :to_h
+      end
+    end
+
+    describe '#metadata' do
+      it 'returns search metadata for each result' do
+        expect(result.metadata).to match(
+          'books_index' => {
+            'indexUid' => 'books_index',
+            'query' => 'Steve', 'processingTimeMs' => 0, 'limit' => 20, 'offset' => 0, 'estimatedTotalHits' => 1
+          },
+          'products_index' => {
+            'indexUid' => 'products_index',
+            'query' => 'palm', 'processingTimeMs' => 0, 'limit' => 1, 'offset' => 0, 'estimatedTotalHits' => 2
+          },
+          'color_index' => {
+            'indexUid' => 'color_index',
+            'query' => 'bl', 'processingTimeMs' => 0, 'limit' => 20, 'offset' => 0, 'estimatedTotalHits' => 2
+          }
+        )
+      end
+    end
+  end
 end