From 1259efb6a83148827445849232dc34e9342c1f96 Mon Sep 17 00:00:00 2001
From: Michael Williamson <mike@zwobble.org>
Date: Sun, 9 Jul 2023 11:50:57 +0100
Subject: Check function returns correct type

---
 .../TypeCheckerFunctionDefinitionTests.java        | 32 ++++++++++++++++++++++
 1 file changed, 32 insertions(+)

(limited to 'src/test/java')

diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
index 25b9c158..7bb490a8 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionDefinitionTests.java
@@ -221,6 +221,38 @@ public class TypeCheckerFunctionDefinitionTests {
         assertThat(error.actual(), equalTo(Types.INT));
     }
 
+    @Test
+    public void givenReturnTypeIsUnitWhenBodyHasNoReturnThenFunctionTypeChecks() {
+        var untyped = UntypedFunctionDefinitionNode.builder()
+            .withReturnType(Untyped.typeLevelReference("Unit"))
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterNamespaceWithDefinedVariables(P.map(
+                Map.entry("Unit", Types.metaType(Types.UNIT))
+            ));
+
+        TypeChecker.typeCheckNamespaceStatement(untyped, context);
+    }
+
+    @Test
+    public void givenReturnTypeIsNotUnitWhenBodyHasNoReturnThenFunctionFailsTypeCheck() {
+        var untyped = UntypedFunctionDefinitionNode.builder()
+            .withReturnType(Untyped.typeLevelReference("String"))
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterNamespaceWithDefinedVariables(P.map(
+                Map.entry("String", Types.metaType(Types.STRING))
+            ));
+
+        var error = assertThrows(
+            UnexpectedTypeError.class,
+            () -> TypeChecker.typeCheckNamespaceStatement(untyped, context)
+        );
+
+        assertThat(error.expected(), equalTo(TypeSets.type(Types.STRING)));
+        assertThat(error.actual(), equalTo(Types.UNIT));
+    }
+
     @Test
     public void freeVariablesIncludeVariablesReferencedInFunctionButDefinedOutsideOfFunction() {
         var untyped = UntypedFunctionDefinitionNode.builder()
-- 
cgit v1.2.3