Merge pull request #324 from Shopify/fix-code-lens

Don't push group id to stack when encountering non-test classes

lib/ruby_lsp/ruby_lsp_rails/code_lens.rb

@@ -88,15 +88,17 @@ def on_class_node_enter(node)
         if class_name.end_with?("Test")
           command = "#{BASE_COMMAND} #{@path}"
           add_test_code_lens(node, name: class_name, command: command, kind: :group)
+          @group_id_stack.push(@group_id)
+          @group_id += 1
         end
-
-        @group_id_stack.push(@group_id)
-        @group_id += 1
       end
 
       sig { params(node: Prism::ClassNode).void }
       def on_class_node_leave(node)
-        @group_id_stack.pop
+        class_name = node.constant_path.slice
+        if class_name.end_with?("Test")
+          @group_id_stack.pop
+        end
       end
 
       private

test/ruby_lsp_rails/code_lens_test.rb

@@ -205,6 +205,60 @@ class NestedTest < ActiveSupport::TestCase
         assert_empty(data)
       end
 
+      test "recognizes nested class structure correctly" do
+        response = generate_code_lens_for_source(<<~RUBY)
+          module Foo
+            class Bar
+              class Test < ActiveSupport::TestCase
+                test "an example" do
+                  # test body
+                end
+              end
+            end
+
+            class AnotherTest < ActiveSupport::TestCase
+              test "an example" do
+                # test body
+              end
+            end
+          end
+        RUBY
+
+        data = response.map(&:data)
+
+        # Code lenses for `Test`
+        explorer, terminal, debug = data.shift(3)
+        assert_nil(explorer[:group_id])
+        assert_nil(terminal[:group_id])
+        assert_nil(debug[:group_id])
+        assert_equal(1, explorer[:id])
+        assert_equal(1, terminal[:id])
+        assert_equal(1, debug[:id])
+
+        # Code lenses for `an example`
+        explorer, terminal, debug = data.shift(3)
+        assert_equal(1, explorer[:group_id])
+        assert_equal(1, terminal[:group_id])
+        assert_equal(1, debug[:group_id])
+
+        # Code lenses for `AnotherTest`
+        explorer, terminal, debug = data.shift(3)
+        assert_nil(explorer[:group_id])
+        assert_nil(terminal[:group_id])
+        assert_nil(debug[:group_id])
+        assert_equal(2, explorer[:id])
+        assert_equal(2, terminal[:id])
+        assert_equal(2, debug[:id])
+
+        # Code lenses for `an example`
+        explorer, terminal, debug = data.shift(3)
+        assert_equal(2, explorer[:group_id])
+        assert_equal(2, terminal[:group_id])
+        assert_equal(2, debug[:group_id])
+
+        assert_empty(data)
+      end
+
       private
 
       def generate_code_lens_for_source(source)