Infer return type of function expressions properly

1 files changed, 2 insertions, 2 deletions

diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionExpressionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionExpressionTests.java
index 80627018..35d61efd 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionExpressionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionExpressionTests.java
@@ -43,7 +43,7 @@ public class TypeCheckerFunctionExpressionTests {
     public void returnTypeCanBeInferred() {
         var untyped = UntypedFunctionExpressionNode.builder()
             .withoutReturnType()
-            .addStatement(Untyped.expressionStatement(Untyped.stringLiteral("")))
+            .addStatement(Untyped.functionReturn(Untyped.stringLiteral("")))
             .build();
         var context = TypeCheckerContext.EMPTY;
 
@@ -52,7 +52,7 @@ public class TypeCheckerFunctionExpressionTests {
         assertThat(
             typed.value(),
             isTypedFunctionExpressionNode()
-                .withType(new FunctionType(List.of(), List.of(), Effects.NONE, Types.UNIT))
+                .withType(new FunctionType(List.of(), List.of(), Effects.NONE, Types.STRING))
         );
     }
 }