Return test

1 files changed, 11 insertions, 2 deletions

diff --git a/test/Language.hs b/test/Language.hs
index ea387d8..a9da577 100644
--- a/test/Language.hs
+++ b/test/Language.hs
@@ -34,9 +34,10 @@ testCase = TestCase $ do
   assertAst
     "parse a function with return value"
     "module main\n\
-    \def fn(): u8 { }"
+    \def fn(): u8 {\n\
+    \return 1; }"
     [ A.Module "main" $ newPos "test" 1 1,
-      A.Func "fn" [] (Just $ A.Type "u8") [] False False $ newPos "test" 2 1
+      A.Func "fn" [] (Just $ A.Type "u8") [A.Return (Just $ A.Num 1 $ newPos "test" 3 8) $ newPos "test" 3 1] False False $ newPos "test" 2 1
     ]
   assertAst
     "parse a function call"
@@ -56,5 +57,13 @@ testCase = TestCase $ do
       A.Func "fn" [("a", A.Type "u8", newPos "test" 2 8)] Nothing [] False False $ newPos "test" 2 1,
       A.Call (A.Var "fn" $ newPos "test" 3 1) [A.Num 10 $ newPos "test" 3 4] $ newPos "test" 3 1
     ]
+  assertAst
+    "parse empty return on a function"
+    "module main\n\
+    \def fn() {\n\
+    \return; }"
+    [ A.Module "main" $ newPos "test" 1 1,
+      A.Func "fn" [] Nothing [A.Return Nothing $ newPos "test" 3 1] False False $ newPos "test" 2 1
+    ]
 
 language = [testCase]