commit 83683ed0aeb6ab38d518837698d90c50b4adf65b
Author: Tomer Brisker <tbrisker@gmail.com>
Date:   Thu Sep 4 10:16:39 2014 +0300

    Fixes #844 #5725 - correct hosts count in puppetclass

diff --git a/app/models/concerns/host_common.rb b/app/models/concerns/host_common.rb
index 05b258ba3..ef2bf6cf5 100644
--- a/app/models/concerns/host_common.rb
+++ b/app/models/concerns/host_common.rb
@@ -203,6 +203,8 @@ module HostCommon
   def update_config_group_counters(record)
     record.update_attribute(:hostgroups_count, cnt_hostgroups(record))
     record.update_attribute(:hosts_count, cnt_hosts(record))
+
+    record.update_puppetclasses_total_hosts
   end
 
 end
diff --git a/app/models/concerns/puppetclass_total_hosts.rb b/app/models/concerns/puppetclass_total_hosts.rb
new file mode 100644
index 000000000..07754c745
--- /dev/null
+++ b/app/models/concerns/puppetclass_total_hosts.rb
@@ -0,0 +1,32 @@
+module PuppetclassTotalHosts
+
+  module Indirect
+    extend ActiveSupport::Concern
+
+    included do
+      #updates counters for all puppetclasses affected indirectly
+      def update_puppetclasses_total_hosts(relation = nil)
+        if self.is_a?(Hostgroup)
+          config_groups.each(&:update_puppetclasses_total_hosts) if config_groups.present?
+          parent.update_puppetclasses_total_hosts unless is_root?
+        end
+        puppetclasses.each(&:update_total_hosts) if puppetclasses.present?
+      end
+
+    end
+  end
+
+  module JoinTable
+    extend ActiveSupport::Concern
+
+    included do
+      after_save :update_total_hosts
+      after_destroy :update_total_hosts
+
+      def update_total_hosts
+        puppetclass.update_total_hosts
+      end
+    end
+  end
+
+end
diff --git a/app/models/config_group.rb b/app/models/config_group.rb
index df7a635cd..e243149ce 100644
--- a/app/models/config_group.rb
+++ b/app/models/config_group.rb
@@ -2,13 +2,14 @@ class ConfigGroup < ActiveRecord::Base
   audited :allow_mass_assignment => true
   include Authorizable
   include Parameterizable::ByIdName
+  include PuppetclassTotalHosts::Indirect
 
   validates_lengths_from_database
 
   attr_accessible :name, :puppetclass_ids
 
   has_many :config_group_classes
-  has_many :puppetclasses, :through => :config_group_classes
+  has_many :puppetclasses, :through => :config_group_classes, :dependent => :destroy
   has_many :host_config_groups
   has_many_hosts :through => :host_config_groups
 
diff --git a/app/models/config_group_class.rb b/app/models/config_group_class.rb
index a33f60dc0..dcab93bde 100644
--- a/app/models/config_group_class.rb
+++ b/app/models/config_group_class.rb
@@ -1,6 +1,7 @@
 class ConfigGroupClass < ActiveRecord::Base
   include Authorizable
   include CounterCacheFix
+  include PuppetclassTotalHosts::JoinTable
 
   audited :associated_with => :config_group, :allow_mass_assignment => true
   attr_accessible :config_group_id, :puppetclass_id
diff --git a/app/models/host/managed.rb b/app/models/host/managed.rb
index 8955bffb3..57b8663ad 100644
--- a/app/models/host/managed.rb
+++ b/app/models/host/managed.rb
@@ -1,11 +1,10 @@
 class Host::Managed < Host::Base
   include ReportCommon
   include Hostext::Search
-
   PROVISION_METHODS = %w[build image]
 
-  has_many :host_classes, :dependent => :destroy, :foreign_key => :host_id
-  has_many :puppetclasses, :through => :host_classes
+  has_many :host_classes, :foreign_key => :host_id
+  has_many :puppetclasses, :through => :host_classes, :dependent => :destroy
   belongs_to :hostgroup
   has_many :reports, :dependent => :destroy, :foreign_key => :host_id
   has_many :host_parameters, :dependent => :destroy, :foreign_key => :reference_id, :inverse_of => :host
@@ -35,6 +34,7 @@ class Host::Managed < Host::Base
   # Custom hooks will be executed after_commit
   after_commit :build_hooks
   before_save :clear_data_on_build
+  after_save :update_hostgroups_puppetclasses, :if => :hostgroup_id_changed?
 
   def build_hooks
     return unless respond_to?(:old) && old && (build? != old.build?)
@@ -958,4 +958,9 @@ class Host::Managed < Host::Base
     LookupValue.where(:match => "fqdn=#{fqdn_was}").update_all(:match => lookup_value_match)
   end
 
+  def update_hostgroups_puppetclasses
+    Hostgroup.find(hostgroup_id_was).update_puppetclasses_total_hosts if hostgroup_id_was.present?
+    Hostgroup.find(hostgroup_id).update_puppetclasses_total_hosts     if hostgroup_id.present?
+  end
+
 end
diff --git a/app/models/host_class.rb b/app/models/host_class.rb
index 4e36a70f1..1c49e326f 100644
--- a/app/models/host_class.rb
+++ b/app/models/host_class.rb
@@ -1,11 +1,11 @@
 class HostClass < ActiveRecord::Base
   include Authorizable
-  include CounterCacheFix
+  include PuppetclassTotalHosts::JoinTable
 
   validates_lengths_from_database
   audited :associated_with => :host, :allow_mass_assignment => true
-  belongs_to_host :foreign_key => :host_id
-  belongs_to :puppetclass, :counter_cache => :hosts_count
+  belongs_to_host
+  belongs_to :puppetclass
 
   validates :host_id, :presence => true
   validates :puppetclass_id, :presence => true, :uniqueness => {:scope => :host_id}
diff --git a/app/models/hostgroup.rb b/app/models/hostgroup.rb
index b32989ba6..0331eec41 100644
--- a/app/models/hostgroup.rb
+++ b/app/models/hostgroup.rb
@@ -6,20 +6,23 @@ class Hostgroup < ActiveRecord::Base
   include HostCommon
   include NestedAncestryCommon
   include ScopedSearchExtensions
+  include PuppetclassTotalHosts::Indirect
 
   validates_lengths_from_database :except => [:name]
   before_destroy EnsureNotUsedBy.new(:hosts)
-  has_many :hostgroup_classes, :dependent => :destroy
-  has_many :puppetclasses, :through => :hostgroup_classes
+  has_many :hostgroup_classes
+  has_many :puppetclasses, :through => :hostgroup_classes, :dependent => :destroy
   validates :name, :format => { :with => /\A(\S+\s?)+\Z/, :message => N_("can't contain trailing white spaces.")}
   validates :root_pass, :allow_blank => true, :length => {:minimum => 8, :message => _('should be 8 characters or more')}
   has_many :group_parameters, :dependent => :destroy, :foreign_key => :reference_id, :inverse_of => :hostgroup
   accepts_nested_attributes_for :group_parameters, :allow_destroy => true
   include ParameterValidators
-  has_many_hosts
+  has_many_hosts :after_add => :update_puppetclasses_total_hosts,
+                 :after_remove => :update_puppetclasses_total_hosts
   has_many :template_combinations, :dependent => :destroy
   has_many :config_templates, :through => :template_combinations
   before_save :remove_duplicated_nested_class
+  after_save :update_ancestry_puppetclasses, :if => :ancestry_changed?
 
   alias_attribute :arch, :architecture
   alias_attribute :os, :operatingsystem
@@ -183,6 +186,11 @@ class Hostgroup < ActiveRecord::Base
     new
   end
 
+  def update_ancestry_puppetclasses
+    unscoped_find(ancestry_was.to_s.split('/').last.to_i).update_puppetclasses_total_hosts if ancestry_was.present?
+    unscoped_find(ancestry.to_s.split('/').last.to_i).update_puppetclasses_total_hosts if ancestry.present?
+  end
+
   private
 
   def lookup_value_match
diff --git a/app/models/hostgroup_class.rb b/app/models/hostgroup_class.rb
index bbcbd6315..47e304945 100644
--- a/app/models/hostgroup_class.rb
+++ b/app/models/hostgroup_class.rb
@@ -1,6 +1,7 @@
 class HostgroupClass < ActiveRecord::Base
   include Authorizable
   include CounterCacheFix
+  include PuppetclassTotalHosts::JoinTable
 
   audited :associated_with => :hostgroup, :allow_mass_assignment => true
   belongs_to :hostgroup
diff --git a/app/models/puppetclass.rb b/app/models/puppetclass.rb
index dcd223e03..a09445dbc 100644
--- a/app/models/puppetclass.rb
+++ b/app/models/puppetclass.rb
@@ -10,12 +10,12 @@ class Puppetclass < ActiveRecord::Base
   has_many :environment_classes, :dependent => :destroy
   has_many :environments, :through => :environment_classes, :uniq => true
   has_and_belongs_to_many :operatingsystems
-  has_many :hostgroup_classes, :dependent => :destroy
-  has_many :hostgroups, :through => :hostgroup_classes
-  has_many :host_classes, :dependent => :destroy
-  has_many_hosts :through => :host_classes
+  has_many :hostgroup_classes
+  has_many :hostgroups, :through => :hostgroup_classes, :dependent => :destroy
+  has_many :host_classes
+  has_many_hosts :through => :host_classes, :dependent => :destroy
   has_many :config_group_classes
-  has_many :config_groups, :through => :config_group_classes
+  has_many :config_groups, :through => :config_group_classes, :dependent => :destroy
 
   has_many :lookup_keys, :inverse_of => :puppetclass, :dependent => :destroy
   accepts_nested_attributes_for :lookup_keys, :reject_if => lambda { |a| a[:key].blank? }, :allow_destroy => true
@@ -34,7 +34,7 @@ class Puppetclass < ActiveRecord::Base
   default_scope lambda { order('puppetclasses.name') }
 
   scoped_search :on => :name, :complete_value => :true
-  scoped_search :on => :hosts_count
+  scoped_search :on => :total_hosts
   scoped_search :on => :global_class_params_count, :rename => :params_count   # Smart Parameters
   scoped_search :on => :lookup_keys_count, :rename => :variables_count        # Smart Variables
   scoped_search :in => :environments, :on => :name, :complete_value => :true, :rename => "environment"
@@ -82,6 +82,32 @@ class Puppetclass < ActiveRecord::Base
     name.gsub(module_name+"::","")
   end
 
+  # return host ids from config groups by type
+  def host_ids_from_config_groups(host_type)
+    ids = config_groups.joins(:host_config_groups)
+                 .where("host_config_groups.host_type='#{host_type}'")
+                 .pluck('host_config_groups.host_id') unless config_group_classes.empty?
+    ids || []
+  end
+
+  def all_hostgroups(with_descendants = true)
+    ids = hostgroup_ids
+    ids += host_ids_from_config_groups('Hostgroup')
+    hgs = Hostgroup.unscoped.where(:id => ids.uniq)
+    hgs.flat_map(&:subtree).uniq if with_descendants
+  end
+
+  def all_hosts
+    ids = host_ids
+    ids += all_hostgroups.flat_map(&:host_ids)
+    ids += host_ids_from_config_groups('Host::Base')
+    Host::Managed.unscoped.where(:id => ids.uniq)
+  end
+
+  def update_total_hosts
+    update_attribute(:total_hosts, all_hosts.count)
+  end
+
   # Populates the rdoc tree with information about all the classes in your modules.
   #   Firstly, we prepare the modules tree
   #   Secondly we run puppetdoc over the modulespath and manifestdir for all environments
diff --git a/app/views/puppetclasses/index.html.erb b/app/views/puppetclasses/index.html.erb
index 2c9b7ab27..4e3b3125a 100644
--- a/app/views/puppetclasses/index.html.erb
+++ b/app/views/puppetclasses/index.html.erb
@@ -6,8 +6,8 @@
   <tr>
     <th><%= sort :name, :as => s_("Puppetclass|Name") %></th>
     <th><%= sort :environment, :as => _("Environments and documentation") %></th>
-    <th><%= _('Host group') %></th>
-    <th><%= sort :hosts_count, :as => _('Hosts'), :default => "DESC" %></th>
+    <th><%= _('Host groups') %></th>
+    <th><%= sort :total_hosts, :as => _('Hosts'), :default => "DESC" %></th>
     <th><%= sort :params_count, :as => _('Parameters'), :default => "DESC" %></th>
     <th><%= sort :variables_count, :as => _('Variables'), :default => "DESC" %></th>
     <th></th>
@@ -21,7 +21,7 @@
         <% end %>
       </td>
       <td><%= puppetclass.hostgroups.map {|hg| link_to_if_authorized trunc(hg), hash_for_edit_hostgroup_path(:id=>hg).merge(:auth_object => hg, :authorizer => @hostgroups_authorizer)}.to_sentence.html_safe %></td>
-      <td> <%= link_to puppetclass.hosts_count, hosts_path(:search => "class = #{puppetclass.name}")%></td>
+      <td> <%= link_to puppetclass.total_hosts, hosts_path(:search => "class = #{puppetclass.name}")%></td>
       <td><%= puppetclass.global_class_params_count %> </td>
       <td><%= puppetclass.lookup_keys_count %> </td>
       <td>
diff --git a/db/migrate/20141109131448_rename_hosts_count_column.rb b/db/migrate/20141109131448_rename_hosts_count_column.rb
new file mode 100644
index 000000000..912b3e8f4
--- /dev/null
+++ b/db/migrate/20141109131448_rename_hosts_count_column.rb
@@ -0,0 +1,10 @@
+class RenameHostsCountColumn < ActiveRecord::Migration
+  #prevent wierdness with rails treating hosts_count as cached counter in some cases
+  def up
+    rename_column :puppetclasses, :hosts_count, :total_hosts
+  end
+
+  def down
+    rename_column :puppetclasses, :total_hosts, :hosts_count
+  end
+end
diff --git a/db/migrate/20141110084848_fix_puppetclass_total_hosts.rb b/db/migrate/20141110084848_fix_puppetclass_total_hosts.rb
new file mode 100644
index 000000000..f90b5f21f
--- /dev/null
+++ b/db/migrate/20141110084848_fix_puppetclass_total_hosts.rb
@@ -0,0 +1,8 @@
+class FixPuppetclassTotalHosts < ActiveRecord::Migration
+  def up
+    Rake::Task['puppet:fix_total_hosts'].invoke
+  end
+
+  def down
+  end
+end
diff --git a/lib/tasks/fix_cached_counters.rake b/lib/tasks/fix_cached_counters.rake
index a9ab3c5fc..2cbf59157 100644
--- a/lib/tasks/fix_cached_counters.rake
+++ b/lib/tasks/fix_cached_counters.rake
@@ -5,7 +5,7 @@ task :fix_cached_counters => :environment do
     cl.all.each{|el| cl.reset_counters(el.id, :hosts, :hostgroups)}
     puts "#{cl} corrected"
   end
-  Puppetclass.all.each{|el| Puppetclass.reset_counters(el.id, :host_classes, :hostgroup_classes, :lookup_keys)}
+  Puppetclass.all.each{|el| Puppetclass.reset_counters(el.id, :hostgroup_classes, :lookup_keys)}
   puts "Puppetclass corrected"
   Model.all.each{|el| Model.reset_counters(el.id, :hosts)}
   puts "Model corrected"
diff --git a/lib/tasks/puppet.rake b/lib/tasks/puppet.rake
index d2f0de1b4..84bc4c9d4 100644
--- a/lib/tasks/puppet.rake
+++ b/lib/tasks/puppet.rake
@@ -217,4 +217,12 @@ namespace :puppet do
     end
   end
 
+  desc "Correct hosts counts for all classes in case they are wrong"
+  task :fix_total_hosts => :environment do
+    if Puppetclass.count > 0
+      User.current = User.anonymous_admin
+      Puppetclass.all.each(&:update_total_hosts)
+    end
+  end
+
 end
diff --git a/test/unit/host_class_test.rb b/test/unit/host_class_test.rb
deleted file mode 100644
index 6d161709a..000000000
--- a/test/unit/host_class_test.rb
+++ /dev/null
@@ -1,19 +0,0 @@
-require 'test_helper'
-
-class HostClassTest < ActiveSupport::TestCase
-
-  setup do
-    disable_orchestration
-    User.current = User.find_by_login "one"
-    # puppetclasses(:two) needs to be in production environment
-    EnvironmentClass.create(:puppetclass_id => puppetclasses(:two).id, :environment_id => environments(:production).id )
-  end
-
-  test "should update hosts_count" do
-    pc = puppetclasses(:two)
-    assert_difference "pc.hosts_count" do
-      hc = HostClass.create(:puppetclass_id => pc.id, :host_id => FactoryGirl.create(:host).id)
-      pc.reload
-    end
-  end
-end
diff --git a/test/unit/puppetclass_test.rb b/test/unit/puppetclass_test.rb
index e29674e22..2a2cc50bc 100644
--- a/test/unit/puppetclass_test.rb
+++ b/test/unit/puppetclass_test.rb
@@ -76,4 +76,134 @@ class PuppetclassTest < ActiveSupport::TestCase
     assert_equal "Puppetclass", "puppetclass".classify
     assert_equal "Puppetclass", "puppetclasses".classify
   end
+
+  context 'host counter updates on all possible class inheritance' do
+    setup do
+      @class = FactoryGirl.create(:puppetclass)
+      @host = FactoryGirl.create(:host)
+      @hostgroup = FactoryGirl.create(:hostgroup)
+      @another_hostgroup = FactoryGirl.create(:hostgroup)
+      @config_group = FactoryGirl.create(:config_group)
+    end
+
+    def check_host_hostgroup
+      assert_difference('@class.total_hosts') do
+        @host.update_attribute(:hostgroup, @hostgroup)
+        @class.reload
+      end
+      assert_difference('@class.total_hosts', -1) do
+        @host.update_attribute(:hostgroup, nil)
+        @class.reload
+      end
+      assert_difference('@class.total_hosts') do
+        @hostgroup.hosts << @host
+        @class.reload
+      end
+      assert_difference('@class.total_hosts', -1) do
+        @hostgroup.hosts.delete(@host)
+        @class.reload
+      end
+    end
+
+    def check_object_class(obj)
+      assert_difference('@class.total_hosts') do
+        obj.puppetclasses << @class
+        @class.reload
+      end
+      assert_difference('@class.total_hosts', -1) do
+        obj.puppetclasses.delete(@class)
+        @class.reload
+      end
+      assert_difference('@class.total_hosts') do
+        @class.send("#{obj.class.table_name}") << obj
+        @class.reload
+      end
+      assert_difference('@class.total_hosts', -1) do
+        @class.send("#{obj.class.table_name}").delete(obj)
+        @class.reload
+      end
+    end
+
+    def check_object_config_group(obj)
+      assert_difference('@class.total_hosts') do
+        obj.config_groups << @config_group
+        @class.reload
+      end
+      assert_difference('@class.total_hosts', -1) do
+        obj.config_groups.delete(@config_group)
+        @class.reload
+      end
+    end
+
+    it 'on direct assignment to host' do
+      check_object_class(@host)
+    end
+
+    # hostgroup-related
+
+    it 'class on hostgroup, adding host to hostgroup' do
+      @hostgroup.puppetclasses << @class
+      check_host_hostgroup
+    end
+
+    it 'host in hostgroup, adding class to hostgroup' do
+      @host.update_attribute(:hostgroup, @hostgroup)
+      check_object_class(@hostgroup)
+    end
+
+    it 'class on hostgroup parent, adding host to hostgroup' do
+      @another_hostgroup.puppetclasses << @class
+      @hostgroup.update_attribute(:parent, @another_hostgroup)
+      check_host_hostgroup
+    end
+
+    it 'hostgroup ancestry change' do
+      @another_hostgroup.puppetclasses << @class
+      @host.update_attribute(:hostgroup, @hostgroup)
+      assert_difference('@class.total_hosts') do
+        @hostgroup.update_attribute(:parent, @another_hostgroup)
+        @class.reload
+      end
+      assert_difference('@class.total_hosts', -1) do
+        @hostgroup.update_attribute(:parent, nil)
+        @class.reload
+      end
+    end
+
+    it 'host in hostgroup, adding class to hostgroup parent' do
+      @hostgroup.update_attribute(:parent, @another_hostgroup)
+      @host.update_attribute(:hostgroup, @hostgroup)
+      check_object_class(@another_hostgroup)
+    end
+
+    # config_group related
+    it 'class on config_group, adding host to config_group' do
+      @config_group.puppetclasses << @class
+      check_object_config_group(@host)
+    end
+
+    it 'host in config_group, adding class to config_group' do
+      @host.config_groups << @config_group
+      check_object_class(@config_group)
+    end
+
+    it 'class on config_group, hostgroup in config_group, adding host to hostgroup' do
+      @config_group.puppetclasses << @class
+      @hostgroup.config_groups << @config_group
+      check_host_hostgroup
+    end
+
+    it 'host in hostgroup, hostgroup in config_group, adding class to config_group' do
+      @hostgroup.config_groups << @config_group
+      @host.update_attribute(:hostgroup, @hostgroup)
+      check_object_class(@config_group)
+    end
+
+    it 'class on config_group, host in hostgroup, adding hostgroup to config_group' do
+      @config_group.puppetclasses << @class
+      @host.update_attribute(:hostgroup, @hostgroup)
+      check_object_config_group(@hostgroup)
+    end
+  end
+
 end
