From 85b0399e3de64ee4b902c74270d35d1baf83e13a Mon Sep 17 00:00:00 2001
From: Andrew Konchin <andry.konchin@gmail.com>
Date: Sun, 28 Jun 2020 03:09:31 +0300
Subject: [PATCH] RSpec/NestedGroups Optimize #on_top_level_describe callback

---
 lib/rubocop/cop/rspec/nested_groups.rb       | 19 +++++++++----------
 spec/rubocop/cop/rspec/nested_groups_spec.rb | 16 ++++++++++++++++
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/lib/rubocop/cop/rspec/nested_groups.rb b/lib/rubocop/cop/rspec/nested_groups.rb
index 87da08764..e626df506 100644
--- a/lib/rubocop/cop/rspec/nested_groups.rb
+++ b/lib/rubocop/cop/rspec/nested_groups.rb
@@ -97,13 +97,11 @@ 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
@@ -111,13 +109,14 @@ def on_top_level_describe(node, _args)
 
         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
 
diff --git a/spec/rubocop/cop/rspec/nested_groups_spec.rb b/spec/rubocop/cop/rspec/nested_groups_spec.rb
index e9d47f617..64640293f 100644
--- a/spec/rubocop/cop/rspec/nested_groups_spec.rb
+++ b/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