Raise error in #scroll_batches when search backend returns a failure

lib/chewy/errors.rb

@@ -36,4 +36,7 @@ def initialize(join_field_type, join_field_name, relations)
       super("`#{join_field_type}` set for the join field `#{join_field_name}` is not on the :relations list (#{relations})")
     end
   end
+
+  class MissingHitsInScrollError < Error
+  end
 end

lib/chewy/search/scrolling.rb

@@ -29,19 +29,26 @@ def scroll_batches(batch_size: Request::DEFAULT_BATCH_SIZE, scroll: Request::DEF
 
         result = perform(size: batch_size, scroll: scroll)
         total = [raw_limit_value, result.fetch('hits', {}).fetch('total', {}).fetch('value', 0)].compact.min
+
+        total_batches = total / batch_size
         last_batch_size = total % batch_size
-        fetched = 0
+
+        total_batches += 1 if last_batch_size != 0
+
         scroll_id = nil
 
-        loop do
+        total_batches.times do |batch_counter|
+          last_run = total_batches - 1 == batch_counter
+
           hits = result.fetch('hits', {}).fetch('hits', [])
-          fetched += hits.size
-          hits = hits.first(last_batch_size) if last_batch_size != 0 && fetched >= total
+          hits = hits.first(last_batch_size) if last_run && last_batch_size != 0
+
+          raise Chewy::MissingHitsInScrollError if hits.empty?
+
           yield(hits) if hits.present?
           scroll_id = result['_scroll_id']
-          break if fetched >= total
 
-          result = perform_scroll(scroll: scroll, scroll_id: scroll_id)
+          result = perform_scroll(scroll: scroll, scroll_id: scroll_id) unless last_run
         end
       ensure
         Chewy.client.clear_scroll(body: {scroll_id: scroll_id}) if scroll_id

spec/chewy/search/scrolling_spec.rb

@@ -33,6 +33,56 @@
     let(:countries) { Array.new(3) { |i| Country.create!(rating: i + 2, name: "country #{i}") } }
 
     describe '#scroll_batches' do
+      describe 'with search backend returning failures' do
+        before do
+          expect(Chewy.client).to receive(:scroll).once.and_return(
+            'hits' => {
+              'total' => {
+                'value' => 5
+              },
+              'hits' => []
+            },
+            '_shards' => {
+              'total' => 5,
+              'successful' => 2,
+              'skipped' => 0,
+              'failed' => 3,
+              'failures' => [
+                {
+                  'shard' => -1,
+                  'index' => nil,
+                  'reason' => {
+                    'type' => 'search_context_missing_exception',
+                    'reason' => 'No search context found for id [34462229]'
+                  }
+                },
+                {
+                  'shard' => -1,
+                  'index' => nil,
+                  'reason' => {
+                    'type' => 'search_context_missing_exception',
+                    'reason' => 'No search context found for id [34462228]'
+                  }
+                },
+                {
+                  'shard' => -1,
+                  'index' => nil,
+                  'reason' => {
+                    'type' => 'search_context_missing_exception',
+                    'reason' => 'No search context found for id [34888662]'
+                  }
+                }
+              ]
+            },
+            '_scroll_id' => 'scroll_id'
+          )
+        end
+
+        specify do
+          expect { request.scroll_batches(batch_size: 2) {} }.to raise_error(Chewy::MissingHitsInScrollError)
+        end
+      end
+
       context do
         before { expect(Chewy.client).to receive(:scroll).twice.and_call_original }
         specify do