Fixed bugs with *_including_tree argument list

lib/closure_tree/has_closure_tree_root.rb

@@ -15,18 +15,16 @@ def has_closure_tree_root(assoc_name, options = {})
       has_one assoc_name, -> { where(parent: nil) }, options
 
       # Fetches the association, eager loading all children and given associations
-      define_method("#{assoc_name}_including_tree") do |assoc_map_or_reload = nil, assoc_map = nil|
+      define_method("#{assoc_name}_including_tree") do |*args|
         reload = false
-        if assoc_map_or_reload.is_a?(::Hash)
-          assoc_map = assoc_map_or_reload
-        else
-          reload = assoc_map_or_reload
-        end
+        reload = args.shift if args && (args.first == true || args.first == false)
+        assoc_map = args
+        assoc_map = [nil] if assoc_map.blank?
 
+        # Memoize
+        @closure_tree_roots ||= {}
+        @closure_tree_roots[assoc_name] ||= {}
         unless reload
-          # Memoize
-          @closure_tree_roots ||= {}
-          @closure_tree_roots[assoc_name] ||= {}
           if @closure_tree_roots[assoc_name].has_key?(assoc_map)
             return @closure_tree_roots[assoc_name][assoc_map]
           end
@@ -52,7 +50,7 @@ def has_closure_tree_root(assoc_name, options = {})
 
         # Fetch all descendants in constant number of queries.
         # This is the last query-triggering statement in the method.
-        temp_root.self_and_descendants.includes(assoc_map).each do |node|
+        temp_root.self_and_descendants.includes(*assoc_map).each do |node|
           id_hash[node.id] = node
           parent_node = id_hash[node[parent_col_id]]
 

spec/db/schema.rb

@@ -71,6 +71,7 @@
   create_table "contracts" do |t|
     t.integer "user_id", :null => false
     t.integer "contract_type_id"
+    t.string "title"
   end
 
   create_table "contract_types" do |t|

spec/has_closure_tree_root_spec.rb

@@ -26,18 +26,36 @@
     user3.children << user5
     user3.children << user6
 
-    user1.contracts.create!(contract_type: ct1)
-    user2.contracts.create!(contract_type: ct1)
-    user3.contracts.create!(contract_type: ct1)
-    user3.contracts.create!(contract_type: ct2)
-    user4.contracts.create!(contract_type: ct2)
-    user5.contracts.create!(contract_type: ct1)
-    user6.contracts.create!(contract_type: ct2)
+    user1.contracts.create!(title: "Contract 1", contract_type: ct1)
+    user2.contracts.create!(title: "Contract 2", contract_type: ct1)
+    user3.contracts.create!(title: "Contract 3", contract_type: ct1)
+    user3.contracts.create!(title: "Contract 4", contract_type: ct2)
+    user4.contracts.create!(title: "Contract 5", contract_type: ct2)
+    user5.contracts.create!(title: "Contract 6", contract_type: ct1)
+    user6.contracts.create!(title: "Contract 7", contract_type: ct2)
   end
 
   context "with basic config" do
     let!(:group) { Group.create!(name: "TheGroup") }
 
+    it "loads all nodes in a constant number of queries" do
+      expect do
+        root = group_reloaded.root_user_including_tree
+        expect(root.children[0].email).to eq "[email protected]"
+        expect(root.children[0].parent.children[1].email).to eq "[email protected]"
+      end.to_not exceed_query_limit(2)
+    end
+
+    it "loads all nodes plus single association in a constant number of queries" do
+      expect do
+        root = group_reloaded.root_user_including_tree(:contracts)
+        expect(root.children[0].email).to eq "[email protected]"
+        expect(root.children[0].parent.children[1].email).to eq "[email protected]"
+        expect(root.children[0].children[0].contracts[0].user.
+          parent.parent.children[1].children[1].contracts[0].title).to eq "Contract 7"
+      end.to_not exceed_query_limit(3)
+    end
+
     it "loads all nodes and associations in a constant number of queries" do
       expect do
         root = group_reloaded.root_user_including_tree(contracts: :contract_type)
@@ -66,6 +84,10 @@
         to eq "[email protected]"
     end
 
+    it "works if true passed on first call" do
+      expect(group_reloaded.root_user_including_tree(true).email).to eq "[email protected]"
+    end
+
     it "eager loads inverse association to group" do
       expect do
         root = group_reloaded.root_user_including_tree