RSpec/NestedGroups Optimize #on_top_level_describe callback

lib/rubocop/cop/rspec/nested_groups.rb

@@ -97,27 +97,26 @@ class NestedGroups < Cop
           "Configuration key `#{DEPRECATED_MAX_KEY}` for #{cop_name} is " \
           'deprecated in favor of `Max`. Please use that instead.'
 
-        def_node_search :find_contexts, ExampleGroups::ALL.block_pattern
-
         def on_top_level_describe(node, _args)
-          find_nested_contexts(node.parent) do |context, nesting|
+          find_nested_example_groups(node.parent) do |example_group, nesting|
             self.max = nesting
             add_offense(
-              context.send_node,
+              example_group.send_node,
               message: message(nesting)
             )
           end
         end
 
         private
 
-        def find_nested_contexts(node, nesting: 1, &block)
-          find_contexts(node) do |nested_context|
-            yield(nested_context, nesting) if nesting > max_nesting
+        def find_nested_example_groups(node, nesting: 1, &block)
+          example_group = example_group?(node)
+          yield node, nesting if example_group && nesting > max_nesting
+
+          next_nesting = example_group ? nesting + 1 : nesting
 
-            nested_context.each_child_node do |child|
-              find_nested_contexts(child, nesting: nesting + 1, &block)
-            end
+          node.each_child_node(:block, :begin) do |child|
+            find_nested_example_groups(child, nesting: next_nesting, &block)
           end
         end
 

spec/rubocop/cop/rspec/nested_groups_spec.rb

@@ -76,4 +76,20 @@ class MyThingy
         ).to_stderr
     end
   end
+
+  it 'counts nesting correctly when non-spec nesting' do
+    expect_offense(<<-RUBY)
+      describe MyClass do
+        context 'when foo' do
+          context 'when bar' do
+            [].each do |i|
+              context 'when baz' do
+              ^^^^^^^^^^^^^^^^^^ Maximum example group nesting exceeded [4/3].
+              end
+            end
+          end
+        end
+      end
+    RUBY
+  end
 end