kkpan11 · pull · Jan 15, 2025
diff --git a/Data-Structures/Tree/BinarySearchTree.js b/Data-Structures/Tree/BinarySearchTree.js
@@ -36,13 +36,13 @@ const Node = (function Node() {
     visit(output = (value) => console.log(value)) {
       // Recursively go left
       if (this.left !== null) {
-        this.left.visit()
+        this.left.visit(output)
       }
       // Print out value
       output(this.value)
       // Recursively go right
       if (this.right !== null) {
-        this.right.visit()
+        this.right.visit(output)
       }
     }
 
@@ -116,20 +116,23 @@ const Tree = (function () {
     }
 
     // Inorder traversal
-    traverse() {
+    traverse(output = (value) => console.log(value)) {
       if (!this.root) {
         // No nodes are there in the tree till now
         return
       }
-      this.root.visit()
+      this.root.visit(output)
     }
 
     // Start by searching the root
     search(val) {
-      const found = this.root.search(val)
-      if (found !== null) {
-        return found.value
+      if (this.root) {
+        const found = this.root.search(val)
+        if (found !== null) {
+          return found.value
+        }
       }
+
       // not found
       return null
     }

diff --git a/Data-Structures/Tree/test/BinarySearchTree.test.js b/Data-Structures/Tree/test/BinarySearchTree.test.js
@@ -0,0 +1,66 @@
+import { Tree } from '../BinarySearchTree.js'
+
+describe('Binary Search Tree', () => {
+  let tree
+
+  beforeEach(() => {
+    tree = new Tree()
+    tree.addValue(10)
+    tree.addValue(5)
+    tree.addValue(15)
+    tree.addValue(3)
+    tree.addValue(8)
+  })
+
+  test('should add values to the tree', () => {
+    tree.addValue(12)
+
+    expect(tree.search(12)).toBe(12)
+    expect(tree.search(5)).toBe(5)
+    expect(tree.search(15)).toBe(15)
+  })
+
+  test('should perform in-order traversal', () => {
+    const values = []
+    const output = (val) => values.push(val)
+    tree.traverse(output)
+    expect(values).toEqual([3, 5, 8, 10, 15])
+  })
+
+  test('should remove leaf nodes correctly', () => {
+    tree.removeValue(5)
+    expect(tree.search(5)).toBeNull()
+  })
+
+  test('should remove nodes with one child correctly', () => {
+    tree.addValue(12)
+    tree.removeValue(15)
+
+    expect(tree.search(15)).toBeNull()
+    expect(tree.search(12)).toBe(12)
+  })
+
+  test('should remove nodes with two children correctly', () => {
+    tree.addValue(18)
+    tree.removeValue(15)
+
+    expect(tree.search(15)).toBeNull()
+    expect(tree.search(18)).toBe(18)
+  })
+
+  test('should return null for non-existent values', () => {
+    expect(tree.search(20)).toBeNull()
+    expect(tree.search(0)).toBeNull()
+  })
+
+  test('should handle removal of root node correctly', () => {
+    tree.removeValue(10)
+    expect(tree.search(10)).toBeNull()
+  })
+
+  test('should handle empty tree gracefully', () => {
+    const newTree = new Tree()
+    newTree.removeValue(22) // Should not throw
+    expect(newTree.search(22)).toBeNull()
+  })
+})