diff --git a/app/Main.hs b/app/Main.hs
index 39f89cc..716fc38 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -10,8 +10,11 @@ import AST
 import Defines
 
 -- Foo = 42
+test0 :: Tree
+test0 = Node "define" (Just (Leaf (Symbol "foo"))) (Just (Leaf (Number 42)))
+
 test1 :: Tree
-test1 = Node "define" (Just (Leaf (Symbol "foo"))) (Just (Leaf (Number 42)))
+test1 = Leaf (Symbol "foo")
 
 -- Bar = 21
 test2 :: Tree
@@ -46,15 +49,4 @@ test9 :: Tree
 test9 = Node "+" (Just (Node "*" (Just (Leaf (Number 2))) (Just (Leaf (Number 5))))) (Just (Node "/" (Just (Leaf (Symbol "foo"))) (Just (Leaf (Number 2)))))
 
 main :: IO ()
-main = computeAllAST (Env [])
-    [   
-        test1,
-        test2,
-        test3,
-        test4,
-        test5,
-        test6,
-        test7,
-        test8,
-        test9
-    ]
+main = putStrLn $ show $ computeAllAST (Env []) [test0, test1, test2, test3, test4, test5, test6, test7, test8, test9]
diff --git a/src/ComputeAST.hs b/src/ComputeAST.hs
index 3e56e80..020920a 100644
--- a/src/ComputeAST.hs
+++ b/src/ComputeAST.hs
@@ -84,30 +84,20 @@ resolveDeepestNode _ _ = (Leaf (Number 0))
 
 ------------ COMPUTE TREE ----------
 
--- Call recursively resolveDeepestNode until the tree is fully resolved (node is a leaf)
-computeTree :: Env -> Tree -> IO Env
--- One node remaining symbol
-computeTree env (Leaf (Symbol symbol)) = do
-    let result = getSymbolValue env symbol
-    putStrLn $ show result
-    return env
--- One node remaining number
-computeTree env (Leaf (Number number)) = do
-    putStrLn $ show number
-    return env
--- More than one node remaining, call resolveDeepestNode on the tree
+computeTree :: Env -> Tree -> Atom
+computeTree env (Leaf (Symbol symbol)) = Number (getSymbolValue env symbol)
+computeTree _ (Leaf (Number number)) = Number number
 computeTree env tree = computeTree env (resolveDeepestNode env tree)
 
 ------------ COMPUTE AST ------------
 
 -- Call appropriate function depending on the node
-computeAST :: Env -> Tree -> IO Env
-computeAST env tree@(Node "define" _ _) = return $ registerDefine env tree
-computeAST env tree = computeTree env tree
+computeAST :: Env -> [Atom] -> Tree -> (Env, [Atom])
+computeAST env atoms tree@(Node "define" _ _) = (registerDefine env tree, atoms)
+computeAST env atoms tree = (env, atoms ++ [computeTree env tree])
 
 -- Call computeAST on every tree in the list
-computeAllAST :: Env -> [Tree] -> IO ()
-computeAllAST _ [] = return ()
-computeAllAST env (tree:rest) = do
-    updatedEnv <- computeAST env tree
-    computeAllAST updatedEnv rest
+computeAllAST :: Env -> [Tree] -> [Atom]
+computeAllAST _ [] = []
+computeAllAST env (tree:rest) = atoms ++ computeAllAST newEnv rest
+    where (newEnv, atoms) = computeAST env [] tree
diff --git a/test/Spec.hs b/test/Spec.hs
index 48a0fce..05da02e 100644
--- a/test/Spec.hs
+++ b/test/Spec.hs
@@ -4,12 +4,14 @@ import Test.Tasty.Runners.Html
 
 import AST
 import TextToAST
+import ComputeAST
+import Defines
 
 main :: IO ()
 main = defaultMainWithIngredients (htmlRunner : defaultIngredients) tests
 
 tests :: TestTree
-tests = testGroup "Tests" [unitTestsASTEqual, unitTestsASTParse]
+tests = testGroup "Tests" [unitTestsASTEqual, unitTestsASTParse, unitTestASTCompute]
 
 unitTestsASTEqual :: TestTree
 unitTestsASTEqual = testGroup "AST Equal Tests"
@@ -150,3 +152,15 @@ unitTestsASTParse = testGroup "AST Parse Tests"
         )
         (textToAST "(foo def )")
   ]
+
+unitTestASTCompute :: TestTree
+unitTestASTCompute = testGroup "AST compute Tests"
+    [ testCase "test1" $
+        assertEqual "number 21 + number 21 = 42"
+          [Number 42]
+          (computeAllAST (Env []) [Node "+" (Just (Leaf (Number 21))) (Just (Leaf (Number 21)))])
+    , testCase "test2" $
+        assertEqual "define foo 42 and tree with leaf foo"
+          [Number 42]
+          (computeAllAST (Env []) [(Node "define" (Just (Leaf (Symbol "foo"))) (Just (Leaf (Number 42)))), (Leaf (Symbol "foo"))])
+    ]