[#Fix 334] Make body consistent for RescueNode, EnsureNode and …

…`KeywordBeginNode`.

changelog/change_fix_334_make_body_consistent_for.md

@@ -0,0 +1 @@
+* [#335](https://github.com/rubocop/rubocop-ast/issues/334): **(Potentially breaking)** [#Fix 334] Make `body` consistent for `RescueNode`, `EnsureNode` and `KeywordBeginNode`. ([@dvandersluis][])

lib/rubocop/ast/node/ensure_node.rb

@@ -6,18 +6,27 @@ module AST
     # node when the builder constructs the AST, making its methods available
     # to all `ensure` nodes within RuboCop.
     class EnsureNode < Node
-      # Returns the body of the `ensure` clause.
+      # Returns the body of the `ensure` node.
       #
-      # @return [Node, nil] The body of the `ensure`.
+      # @return [Node, nil] The body of the `ensure` node.
       def body
+        return rescue_node.body if rescue_node
+
+        node_parts[0]
+      end
+
+      # Returns an the ensure branch in the exception handling statement.
+      #
+      # @return [Node, nil] the body of the ensure branch.
+      def branch
         node_parts[1]
       end
 
       # Returns the `rescue` node of the `ensure`, if present.
       #
       # @return [Node, nil] The `rescue` node.
       def rescue_node
-        node_parts[0] if node_parts[0].rescue_type?
+        node_parts[0] if node_parts[0]&.rescue_type?
       end
 
       # Checks whether this node body is a void context.

lib/rubocop/ast/node/keyword_begin_node.rb

@@ -16,7 +16,7 @@ def body
         if rescue_node
           rescue_node.body
         elsif ensure_node
-          ensure_node.node_parts[0]
+          ensure_node.body
         elsif node_parts.one?
           node_parts[0]
         else

spec/rubocop/ast/ensure_node_spec.rb

@@ -1,7 +1,9 @@
 # frozen_string_literal: true
 
 RSpec.describe RuboCop::AST::EnsureNode do
-  let(:ensure_node) { parse_source(source).ast.children.first }
+  let(:parsed_source) { parse_source(source) }
+  let(:ensure_node) { parsed_source.ast.children.first }
+  let(:node) { parsed_source.node }
 
   describe '.new' do
     let(:source) { 'begin; beginbody; ensure; ensurebody; end' }
@@ -10,9 +12,57 @@
   end
 
   describe '#body' do
-    let(:source) { 'begin; beginbody; ensure; :ensurebody; end' }
+    subject(:body) { ensure_node.body }
 
-    it { expect(ensure_node.body).to be_sym_type }
+    context 'when there is no body' do
+      let(:source) { 'begin; ensure; ensurebody; end' }
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'when the body is a single line' do
+      let(:source) { 'begin; >>beginbody<<; ensure; ensurebody; end' }
+
+      it { is_expected.to eq(node) }
+    end
+
+    context 'when the body is multiple lines' do
+      let(:source) { 'begin; >>foo<<; bar; ensure; ensurebody; end' }
+
+      it 'returns a begin node' do
+        expect(body).to be_begin_type
+        expect(body.children).to include(node)
+      end
+    end
+
+    context 'with `rescue`' do
+      context 'when there is no body' do
+        let(:source) { 'begin; rescue; rescuebody; ensure; ensurebody; end' }
+
+        it { is_expected.to be_nil }
+      end
+
+      context 'when the body is a single line' do
+        let(:source) { 'begin; >>beginbody<<; rescue; rescuebody; ensure; ensurebody; end' }
+
+        it { is_expected.to eq(node) }
+      end
+
+      context 'when the body is multiple lines' do
+        let(:source) { 'begin; >>foo<<; bar; rescue; rescuebody; ensure; ensurebody; end' }
+
+        it 'returns a begin node' do
+          expect(body).to be_begin_type
+          expect(body.children).to include(node)
+        end
+      end
+    end
+  end
+
+  describe '#branch' do
+    let(:source) { 'begin; beginbody; ensure; >>ensurebody<<; end' }
+
+    it { expect(ensure_node.branch).to eq(node) }
   end
 
   describe '#rescue_node' do

spec/rubocop/ast/rescue_node_spec.rb

@@ -1,8 +1,10 @@
 # frozen_string_literal: true
 
 RSpec.describe RuboCop::AST::RescueNode do
-  subject(:ast) { parse_source(source).ast }
+  subject(:ast) { parsed_source.ast }
 
+  let(:parsed_source) { parse_source(source) }
+  let(:node) { parsed_source.node }
   let(:rescue_node) { ast.children.first }
 
   describe '.new' do
@@ -16,26 +18,43 @@
   end
 
   describe '#body' do
-    let(:source) { <<~RUBY }
-      begin
-        foo
-      rescue => e
-      end
-    RUBY
+    subject(:body) { rescue_node.body }
 
-    it { expect(rescue_node.body).to be_send_type }
+    context 'when the body is empty' do
+      let(:source) { <<~RUBY }
+        begin
+        rescue => e
+        end
+      RUBY
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'when the body is a single line' do
+      let(:source) { <<~RUBY }
+        begin
+          >>foo<<
+        rescue => e
+        end
+      RUBY
+
+      it { is_expected.to eq(node) }
+    end
 
     context 'with multiple lines in body' do
       let(:source) { <<~RUBY }
         begin
-          foo
+          >>foo<<
           bar
         rescue => e
           baz
         end
       RUBY
 
-      it { expect(rescue_node.body).to be_begin_type }
+      it 'returns a begin node' do
+        expect(body).to be_begin_type
+        expect(body.children).to include(node)
+      end
     end
   end