Fix ensure_first_param fix method to retrieve the full value of the ensure parameter

lib/puppet-lint/lexer/token.rb

@@ -96,6 +96,80 @@ def to_manifest
           @value
         end
       end
+
+      # Public: Search from this token to find the next token of a given type.
+      #
+      # type - A Symbol type of the token to find, or an Array of Symbols.
+      # opts - An optional Hash
+      #   :value       - A token value to search for in addition to type
+      #   :skip_blocks - A Boolean to specify whether { } blocks should be
+      #                  skipped over (defaults to true).
+      #
+      # Returns a PuppetLint::Lexer::Token object if a matching token could be
+      # found, otherwise nil.
+      def next_token_of(type, opts = {})
+        find_token_of(:next, type, opts)
+      end
+
+      # Public: Search from this token to find the previous token of a given type.
+      #
+      # type - A Symbol type of the token to find, or an Array of Symbols.
+      # opts - An optional Hash
+      #   :value       - A token value to search for in addition to type
+      #   :skip_blocks - A Boolean to specify whether { } blocks should be
+      #                  skipped over (defaults to true).
+      #
+      # Returns a PuppetLint::Lexer::Token object if a matching token could be
+      # found, otherwise nil.
+      def prev_token_of(type, opts = {})
+        find_token_of(:prev, type, opts)
+      end
+
+      # Internal: Search from this token to find the next token of a given type
+      # in a given direction.
+      #
+      # direction - A Symbol direction to search (:next or :prev).
+      # type      - A Symbol type of the token to find, or an Array of Symbols.
+      # opts      - An optional Hash
+      #   :value       - A token value to search for in addition to type
+      #   :skip_blocks - A Boolean to specify whether { } blocks should be
+      #                  skipped over (defaults to true).
+      #
+      # Returns a PuppetLint::Lexer::Token object if a matching token could be
+      # found, otherwise nil.
+      def find_token_of(direction, type, opts = {})
+        return nil unless [:next, :prev].include?(direction)
+
+        opts[:skip_blocks] ||= true
+        to_find = Array[*type]
+
+        token_iter = self.send("#{direction}_token".to_sym)
+        while !token_iter.nil?
+          if to_find.include? token_iter.type
+            if opts[:value]
+              return token_iter if token_iter.value == opts[:value]
+            else
+              return token_iter
+            end
+          end
+
+          opening_token = direction == :next ? "L" : "R"
+          closing_token = direction == :next ? "R" : "L"
+
+          if opts[:skip_blocks]
+            case token_iter.type
+            when "#{opening_token}BRACE".to_sym
+              token_iter = token_iter.send("#{direction}_token_of".to_sym, ["#{closing_token}BRACE".to_sym, opts])
+            when "#{opening_token}BRACK".to_sym
+              token_iter = token_iter.send("#{direction}_token_of".to_sym, ["#{closing_token}BRACK".to_sym, opts])
+            when "#{opening_token}PAREN".to_sym
+              token_iter = token_iter.send("#{direction}_token_of".to_sym, ["#{closing_token}PAREN".to_sym, opts])
+            end
+          end
+          token_iter = token_iter.send("#{direction}_token".to_sym)
+        end
+        nil
+      end
     end
   end
 end

lib/puppet-lint/plugins/check_resources.rb

@@ -47,43 +47,22 @@ def check
       end
     end
   end
+
   def fix(problem)
-    # We find the first and ensure paramss boundaries
-    first_param_name_token = nil
-    first_param_name_idx = nil
-    first_param_comma_token = nil
-    first_param_comma_idx = nil
-    ensure_param_name_token = nil
-    ensure_param_name_idx = nil
-    ensure_param_comma_token = nil
-    ensure_param_comma_idx = nil
-    tokens[(problem[:resource][:start])..(problem[:resource][:end])].each_with_index do |token, token_idx|
-      if first_param_name_token.nil?
-        if token.type == :NAME
-          first_param_name_token = token
-          first_param_name_idx = problem[:resource][:start] + token_idx
-        end
-      elsif first_param_comma_token.nil?
-        if token.type == :COMMA
-          first_param_comma_token = token
-          first_param_comma_idx = problem[:resource][:start] + token_idx
-        end
-      elsif ensure_param_name_token.nil?
-        if token.type == :NAME and token.value == 'ensure'
-          ensure_param_name_token = token
-          ensure_param_name_idx = problem[:resource][:start] + token_idx
-        end
-      elsif ensure_param_comma_token.nil?
-        if token.type == :COMMA or token.type == :SEMIC
-          ensure_param_comma_token = token
-          ensure_param_comma_idx = problem[:resource][:start] + token_idx
-          break
-        end
-      end
-    end
+    first_param_name_token = tokens[problem[:resource][:start]].next_token_of(:NAME)
+    first_param_comma_token = first_param_name_token.next_token_of(:COMMA)
+    ensure_param_name_token = first_param_comma_token.next_token_of(:NAME, :value => 'ensure')
+    ensure_param_comma_token = ensure_param_name_token.next_token_of([:COMMA, :SEMIC])
+
     if first_param_name_token.nil? or first_param_comma_token.nil? or ensure_param_name_token.nil? or ensure_param_comma_token.nil?
       raise PuppetLint::NoFix
     end
+
+    first_param_name_idx = tokens.index(first_param_name_token)
+    first_param_comma_idx = tokens.index(first_param_comma_token)
+    ensure_param_name_idx = tokens.index(ensure_param_name_token)
+    ensure_param_comma_idx = tokens.index(ensure_param_comma_token)
+
     # Flip params
     prev_token = first_param_name_token.prev_token
     first_param_name_token.prev_token = ensure_param_name_token.prev_token

spec/puppet-lint/plugins/check_resources/ensure_first_param_spec.rb

@@ -159,5 +159,45 @@ class {'thang':
         expect(problems).to have(0).problems
       end
     end
+
+    context 'ensure is a selector' do
+      let(:code) { "
+        file { 'foo':
+          mode   => '0640',
+          ensure => $::operatingsystem ? {
+            'redhat' => absent,
+            default  => $::phase_of_the_moon ? {
+              'full'  => absent,
+              default => present,
+            },
+          },
+        }
+      " }
+
+      let(:fixed) { "
+        file { 'foo':
+          ensure => $::operatingsystem ? {
+            'redhat' => absent,
+            default  => $::phase_of_the_moon ? {
+              'full'  => absent,
+              default => present,
+            },
+          },
+          mode   => '0640',
+        }
+      " }
+
+      it 'should detect a problem' do
+        expect(problems).to have(1).problem
+      end
+
+      it 'should fix the problem' do
+        expect(problems).to contain_fixed(msg).on_line(4).in_column(11)
+      end
+
+      it 'should move the whole ensure parameter to the top' do
+        expect(manifest).to eq(fixed)
+      end
+    end
   end
 end