diff --git a/DIRECTORY.md b/DIRECTORY.md
index 1fc9ec0ce9cf..0f7184cdb7e5 100644
--- a/DIRECTORY.md
+++ b/DIRECTORY.md
@@ -912,6 +912,7 @@
               * [StackArrayTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/stacks/StackArrayTest.java)
               * [StackOfLinkedListTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/stacks/StackOfLinkedListTest.java)
             * trees
+              * [AVLTreeTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/trees/AVLTreeTest.java)
               * [BinaryTreeTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/trees/BinaryTreeTest.java)
               * [BoundaryTraversalTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/trees/BoundaryTraversalTest.java)
               * [BSTFromSortedArrayTest](https://github.com/TheAlgorithms/Java/blob/master/src/test/java/com/thealgorithms/datastructures/trees/BSTFromSortedArrayTest.java)
diff --git a/src/main/java/com/thealgorithms/datastructures/trees/AVLTree.java b/src/main/java/com/thealgorithms/datastructures/trees/AVLTree.java
index 7b959b085353..77ee5d5fa23e 100644
--- a/src/main/java/com/thealgorithms/datastructures/trees/AVLTree.java
+++ b/src/main/java/com/thealgorithms/datastructures/trees/AVLTree.java
@@ -1,11 +1,19 @@
 package com.thealgorithms.datastructures.trees;
 
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Represents an AVL Tree, a self-balancing binary search tree.
+ * In an AVL tree, the heights of the two child subtrees of any node
+ * differ by at most one. If they differ by more than one at any time,
+ * rebalancing is performed to restore this property.
+ */
 public class AVLTree {
 
     private Node root;
 
-    private class Node {
-
+    private static class Node {
         private int key;
         private int balance;
         private int height;
@@ -17,8 +25,18 @@ private class Node {
             key = k;
             parent = p;
         }
+
+        public Integer getBalance() {
+            return balance;
+        }
     }
 
+    /**
+     * Inserts a new key into the AVL tree.
+     *
+     * @param key the key to be inserted
+     * @return {@code true} if the key was inserted, {@code false} if the key already exists
+     */
     public boolean insert(int key) {
         if (root == null) {
             root = new Node(key, null);
@@ -31,7 +49,6 @@ public boolean insert(int key) {
                 }
 
                 parent = n;
-
                 boolean goLeft = n.key > key;
                 n = goLeft ? n.left : n.right;
 
@@ -49,8 +66,32 @@ public boolean insert(int key) {
         return true;
     }
 
+    /**
+     * Deletes a key from the AVL tree.
+     *
+     * @param delKey the key to be deleted
+     */
+    public void delete(int delKey) {
+        if (root == null) {
+            return;
+        }
+
+        // Find the node to be deleted
+        Node node = root;
+        Node child = root;
+        while (child != null) {
+            node = child;
+            child = delKey >= node.key ? node.right : node.left;
+            if (delKey == node.key) {
+                delete(node);
+                return;
+            }
+        }
+    }
+
     private void delete(Node node) {
         if (node.left == null && node.right == null) {
+            // Leaf node
             if (node.parent == null) {
                 root = null;
             } else {
@@ -64,6 +105,8 @@ private void delete(Node node) {
             }
             return;
         }
+
+        // Node has one or two children
         Node child;
         if (node.left != null) {
             child = node.left;
@@ -80,26 +123,49 @@ private void delete(Node node) {
         delete(child);
     }
 
-    public void delete(int delKey) {
-        if (root == null) {
-            return;
+    /**
+     * Returns a list of balance factors for each node in the tree.
+     *
+     * @return a list of integers representing the balance factors of the nodes
+     */
+    public List<Integer> returnBalance() {
+        List<Integer> balances = new ArrayList<>();
+        returnBalance(root, balances);
+        return balances;
+    }
+
+    private void returnBalance(Node n, List<Integer> balances) {
+        if (n != null) {
+            returnBalance(n.left, balances);
+            balances.add(n.getBalance());
+            returnBalance(n.right, balances);
         }
-        Node node = root;
-        Node child = root;
+    }
 
-        while (child != null) {
-            node = child;
-            child = delKey >= node.key ? node.right : node.left;
-            if (delKey == node.key) {
-                delete(node);
-                return;
-            }
+    /**
+     * Searches for a key in the AVL tree.
+     *
+     * @param key the key to be searched
+     * @return true if the key is found, false otherwise
+     */
+    public boolean search(int key) {
+        Node result = searchHelper(this.root, key);
+        return result != null;
+    }
+
+    private Node searchHelper(Node root, int key) {
+        if (root == null || root.key == key) {
+            return root;
         }
+
+        if (root.key > key) {
+            return searchHelper(root.left, key);
+        }
+        return searchHelper(root.right, key);
     }
 
     private void rebalance(Node n) {
         setBalance(n);
-
         if (n.balance == -2) {
             if (height(n.left.left) >= height(n.left.right)) {
                 n = rotateRight(n);
@@ -143,7 +209,6 @@ private Node rotateLeft(Node a) {
         }
 
         setBalance(a, b);
-
         return b;
     }
 
@@ -169,7 +234,6 @@ private Node rotateRight(Node a) {
         }
 
         setBalance(a, b);
-
         return b;
     }
 
@@ -197,53 +261,9 @@ private void setBalance(Node... nodes) {
         }
     }
 
-    public void printBalance() {
-        printBalance(root);
-    }
-
-    private void printBalance(Node n) {
-        if (n != null) {
-            printBalance(n.left);
-            System.out.printf("%s ", n.balance);
-            printBalance(n.right);
-        }
-    }
-
     private void reheight(Node node) {
         if (node != null) {
             node.height = 1 + Math.max(height(node.left), height(node.right));
         }
     }
-
-    public boolean search(int key) {
-        Node result = searchHelper(this.root, key);
-        return result != null;
-    }
-
-    private Node searchHelper(Node root, int key) {
-        // root is null or key is present at root
-        if (root == null || root.key == key) {
-            return root;
-        }
-
-        // key is greater than root's key
-        if (root.key > key) {
-            return searchHelper(root.left, key); // call the function on the node's left child
-        }
-        // key is less than root's key then
-        // call the function on the node's right child as it is greater
-        return searchHelper(root.right, key);
-    }
-
-    public static void main(String[] args) {
-        AVLTree tree = new AVLTree();
-
-        System.out.println("Inserting values 1 to 10");
-        for (int i = 1; i < 10; i++) {
-            tree.insert(i);
-        }
-
-        System.out.print("Printing balance: ");
-        tree.printBalance();
-    }
 }
diff --git a/src/test/java/com/thealgorithms/datastructures/trees/AVLTreeTest.java b/src/test/java/com/thealgorithms/datastructures/trees/AVLTreeTest.java
new file mode 100644
index 000000000000..6aa5dc9e22ed
--- /dev/null
+++ b/src/test/java/com/thealgorithms/datastructures/trees/AVLTreeTest.java
@@ -0,0 +1,101 @@
+package com.thealgorithms.datastructures.trees;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.List;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class AVLTreeTest {
+    private AVLTree avlTree;
+
+    @BeforeEach
+    public void setUp() {
+        avlTree = new AVLTree();
+    }
+
+    @Test
+    public void testInsert() {
+        assertTrue(avlTree.insert(10));
+        assertTrue(avlTree.insert(20));
+        assertTrue(avlTree.insert(5));
+        assertFalse(avlTree.insert(10)); // Duplicate
+    }
+
+    @Test
+    public void testSearch() {
+        avlTree.insert(15);
+        avlTree.insert(25);
+        assertTrue(avlTree.search(15));
+        assertFalse(avlTree.search(30)); // Not in the tree
+    }
+
+    @Test
+    public void testDeleteLeafNode() {
+        avlTree.insert(10);
+        avlTree.insert(20);
+        avlTree.insert(30);
+        avlTree.delete(30);
+        assertFalse(avlTree.search(30));
+    }
+
+    @Test
+    public void testDeleteNodeWithOneChild() {
+        avlTree.insert(20);
+        avlTree.insert(10);
+        avlTree.insert(30);
+        avlTree.delete(10);
+        assertFalse(avlTree.search(10));
+    }
+
+    @Test
+    public void testDeleteNodeWithTwoChildren() {
+        avlTree.insert(20);
+        avlTree.insert(10);
+        avlTree.insert(30);
+        avlTree.insert(25);
+        avlTree.delete(20);
+        assertFalse(avlTree.search(20));
+        assertTrue(avlTree.search(30));
+        assertTrue(avlTree.search(25));
+    }
+
+    @Test
+    public void testReturnBalance() {
+        avlTree.insert(10);
+        avlTree.insert(20);
+        avlTree.insert(5);
+        List<Integer> balances = avlTree.returnBalance();
+        assertEquals(3, balances.size()); // There should be 3 nodes
+        assertEquals(0, balances.get(0)); // Balance for node 5
+        assertEquals(0, balances.get(1)); // Balance for node 10
+        assertEquals(0, balances.get(2)); // Balance for node 20
+    }
+
+    @Test
+    public void testInsertAndRebalance() {
+        avlTree.insert(30);
+        avlTree.insert(20);
+        avlTree.insert(10); // This should cause a right rotation
+        assertTrue(avlTree.search(20));
+        assertTrue(avlTree.search(10));
+        assertTrue(avlTree.search(30));
+    }
+
+    @Test
+    public void testComplexInsertionAndDeletion() {
+        avlTree.insert(30);
+        avlTree.insert(20);
+        avlTree.insert(10);
+        avlTree.insert(25);
+        avlTree.insert(5);
+        avlTree.insert(15);
+
+        avlTree.delete(20); // Test deletion
+        assertFalse(avlTree.search(20));
+        assertTrue(avlTree.search(30));
+        assertTrue(avlTree.search(25));
+    }
+}