diff --git a/src/main.c b/src/main.c
index 4e1046e..be31111 100644
--- a/src/main.c
+++ b/src/main.c
@@ -3,6 +3,20 @@
 
 #include "tree.h"
 
+#define ANSI_COLOR_GREEN   "\x1b[32m"
+#define ANSI_COLOR_RED     "\x1b[31m"
+#define ANSI_COLOR_RESET   "\x1b[0m"
+
+int assert(bool condition, const char *message) {
+    if (!condition) {
+        printf(ANSI_COLOR_RED "Assertion failed: %s\n" ANSI_COLOR_RESET, message);
+        exit(1);
+    } else {
+        printf(ANSI_COLOR_GREEN "Assertion passed: %s\n" ANSI_COLOR_RESET, message);
+    }
+    return 0;
+}
+
 int main(void) {
   Tree tree = {NULL};
 
@@ -10,23 +24,28 @@ int main(void) {
   tree_insert(&tree, 1);
   tree_insert(&tree, 2);
   tree_insert(&tree, 3);
-
-  // Check size and print inorder
-  printf("Tree size: %d\n", tree_size(tree.root));
   tree_inorder(tree.root);
+  printf("\n");
 
-  printf("Tree contains 2: %d\n", tree_contains(tree.root, 2));
+  // Basic sanity checks
+  assert(tree.root->data == 1, "Root should be 1");
+  assert(tree.root->right->data == 2, "Right should be 2");
+  assert(tree.root->right->right->data == 3, "Right right should be 3");
+  assert(tree_size(tree.root) == 3, "Tree size should be 3");
+
+  // Test removal
+  assert(tree_contains(tree.root, 2), "Tree should contain 2");
+  tree_remove(&tree.root, 2);
+  assert(!tree_contains(tree.root, 2), "Tree should not contain 2");
 
   // Clear the tree
   tree_clear(&tree);
 
-  // Check if the tree is really empty
-  if (tree.root == NULL)
-    printf("Tree is empty\n");
-  else
-    printf("Tree is not empty: %d\n", tree.root->data);
+  // Check if the tree is empty
+  assert(tree.root == NULL, "Tree should be empty");
+  assert(tree_size(tree.root) == 0, "Tree size should be 0");
 
-  // Check size and print inorder
+  // Print size and inorder traversal
   printf("Tree size: %d\n", tree_size(tree.root));
   tree_inorder(tree.root);