Tree: let the width parameter also be a function that returns an int

2 files changed, 29 insertions, 4 deletions

diff --git a/lib/tree.js b/lib/tree.js
index 75e2d75f..839b698c 100644
--- a/lib/tree.js
+++ b/lib/tree.js
@@ -25,6 +25,7 @@ var tree = {
             };
         }
 
+
         if (width <= 0) {
             throw {
                 name: "TreeParamError",
@@ -46,9 +47,19 @@ var tree = {
                 else {
                     if (depth !== 0) {
                         value = [];
-                        for (var i = 0; i < width; i++) {
+                        var evalWidth = 1;
+
+                        if (typeof(width) == "function") {
+                            evalWidth = width();
+                        }
+                        else {
+                            evalWidth = width;
+                        }
+
+                        for (var i = 0; i < evalWidth; i++) {
                             value.push(this.createTree(depth - 1, width, obj));
                         }
+
                     }
                 }
 
@@ -61,4 +72,4 @@ var tree = {
 
 };
 
-module.exports = tree;
\ No newline at end of file
+module.exports = tree;
diff --git a/test/tree.unit.js b/test/tree.unit.js
index dc6c9530..3af144eb 100644
--- a/test/tree.unit.js
+++ b/test/tree.unit.js
@@ -91,6 +91,20 @@ describe("tree.js", function () {
             Faker.random.first_name.restore();
         });
 
-    });
+        it("can accept a function for the width", function () {
+            var widthFuncCalled = 0;
+            var widthFunc = function () {
+                widthFuncCalled = widthFuncCalled + 1;
+                return 2;
+            };
+
+            var tree = Faker.Tree.createTree(2, widthFunc, proto);
+            assert.equal(widthFuncCalled, 3);
+
+            Faker.random.first_name.restore();
 
-});
+        });
+
+
+    });
+});
\ No newline at end of file