Refactor CheckModelAttributes

lib/brakeman/checks/check_model_attributes.rb

@@ -11,26 +11,35 @@ class Brakeman::CheckModelAttributes < Brakeman::BaseCheck
   def run_check
     return if mass_assign_disabled?
 
-    names = []
 
-    tracker.models.each do |name, model|
-      if model[:attr_accessible].nil? and parent? model, :"ActiveRecord::Base"
-        if tracker.options[:collapse_mass_assignment]
-          names << name.to_s
-        else
-          warn :model => name, 
-            :warning_type => "Attribute Restriction",
-            :message => "Mass assignment is not restricted using attr_accessible", 
-            :confidence => CONFIDENCE[:high]
-        end
+    if tracker.options[:collapse_mass_assignment]
+      names = []
+
+      check_models do |name|
+        names << name.to_s
+      end
+
+      unless names.empty?
+        warn :model => names.sort.join(", "),
+          :warning_type => "Attribute Restriction",
+          :message => "Mass assignment is not restricted using attr_accessible",
+          :confidence => CONFIDENCE[:high]
+      end
+    else
+      check_models do |name|
+        warn :model => name,
+          :warning_type => "Attribute Restriction",
+          :message => "Mass assignment is not restricted using attr_accessible",
+          :confidence => CONFIDENCE[:high]
       end
     end
+  end
 
-    if tracker.options[:collapse_mass_assignment] and not names.empty?
-      warn :model => names.sort.join(", "), 
-        :warning_type => "Attribute Restriction", 
-        :message => "Mass assignment is not restricted using attr_accessible", 
-        :confidence => CONFIDENCE[:high]
+  def check_models
+    tracker.models.each do |name, model|
+      if model[:attr_accessible].nil? and parent? model, :"ActiveRecord::Base"
+        yield name
+      end
     end
   end
 end