diff --git a/src/main/java/com/thealgorithms/datastructures/trees/VerticalOrderTraversal.java b/src/main/java/com/thealgorithms/datastructures/trees/VerticalOrderTraversal.java
index f90fab52..15fd348e 100644
--- a/src/main/java/com/thealgorithms/datastructures/trees/VerticalOrderTraversal.java
+++ b/src/main/java/com/thealgorithms/datastructures/trees/VerticalOrderTraversal.java
@@ -22,23 +22,13 @@ in a tree from top to bottom and left to right, so for a tree :
  */
 public class VerticalOrderTraversal {
 
-    public static void main(String[] args) {
-        BinaryTree tree = new BinaryTree();
-        tree.put(5);
-        tree.put(6);
-        tree.put(3);
-        tree.put(1);
-        tree.put(4);
-        BinaryTree.Node root = tree.getRoot();
-        ArrayList<Integer> ans = verticalTraversal(root);
-        for (int i : ans) {
-            System.out.print(i + " ");
-        }
-    }
-
-    /*Function that receives a root Node and prints the tree 
+    /*Function that receives a root Node and prints the tree
 	in Vertical Order.*/
-    private static ArrayList<Integer> verticalTraversal(BinaryTree.Node root) {
+    public static ArrayList<Integer> verticalTraversal(BinaryTree.Node root) {
+        if (root == null) {
+            return new ArrayList<>();
+        }
+
         /*Queue to store the Nodes.*/
         Queue<BinaryTree.Node> queue = new LinkedList<>();
 
@@ -84,21 +74,19 @@ public class VerticalOrderTraversal {
 			 to the respective ArrayList present at that
 			 index. */
             map.get(index.peek()).add(queue.peek().data);
-            max = (int) Math.max(max, index.peek());
-            min = (int) Math.min(min, index.peek());
-            /*The Node and its index are removed 
+            max = Math.max(max, index.peek());
+            min = Math.min(min, index.peek());
+            /*The Node and its index are removed
 			 from their respective queues.*/
             index.poll();
             queue.poll();
         }
         /*Finally map data is printed here which has keys
-		from min to max. Each ArrayList represents a 
+		from min to max. Each ArrayList represents a
 		vertical column that is added in ans ArrayList.*/
         ArrayList<Integer> ans = new ArrayList<>();
         for (int i = min; i <= max; i++) {
-            for (int j = 0; j < map.get(i).size(); j++) {
-                ans.add(map.get(i).get(j));
-            }
+            ans.addAll(map.get(i));
         }
         return ans;
     }
diff --git a/src/test/java/com/thealgorithms/datastructures/trees/VerticalOrderTraversalTest.java b/src/test/java/com/thealgorithms/datastructures/trees/VerticalOrderTraversalTest.java
new file mode 100644
index 00000000..42875e4f
--- /dev/null
+++ b/src/test/java/com/thealgorithms/datastructures/trees/VerticalOrderTraversalTest.java
@@ -0,0 +1,53 @@
+package com.thealgorithms.datastructures.trees;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+/**
+ * @author Albina Gimaletdinova on 13/01/2023
+ */
+public class VerticalOrderTraversalTest {
+    @Test
+    public void testRootNull() {
+        assertEquals(Collections.emptyList(), VerticalOrderTraversal.verticalTraversal(null));
+    }
+
+    @Test
+    public void testSingleNodeTree() {
+        final BinaryTree.Node root = TreeTestUtils.createTree(new Integer[]{50});
+        assertEquals(List.of(50), VerticalOrderTraversal.verticalTraversal(root));
+    }
+
+    /*
+         1
+        / \
+       2   3
+      /\   /\
+     4  5 6  7
+    */
+    @Test
+    public void testVerticalTraversalCompleteTree() {
+        final BinaryTree.Node root = TreeTestUtils.createTree(new Integer[]{1, 2, 3, 4, 5, 6, 7});
+        assertEquals(List.of(4, 2, 1, 5, 6, 3, 7), VerticalOrderTraversal.verticalTraversal(root));
+    }
+
+    /*
+         1
+        / \
+       2   3
+      /\   /\
+     4   56  7
+        /  \
+       8    9
+    */
+    @Test
+    public void testVerticalTraversalDifferentHeight() {
+        final BinaryTree.Node root = TreeTestUtils.createTree(
+                new Integer[]{1, 2, 3, 4, 5, 6, 7, null, null, 8, null, null, 9});
+        assertEquals(List.of(4, 2, 8, 1, 5, 6, 3, 9, 7), VerticalOrderTraversal.verticalTraversal(root));
+    }
+}