1 files changed, 54 insertions, 0 deletions
diff --git a/src/test/java/org/shedlang/compiler/types/TypeSubstitutionTests.java b/src/test/java/org/shedlang/compiler/types/TypeSubstitutionTests.java
index b7eaa0aa..cd676d71 100644
--- a/src/test/java/org/shedlang/compiler/types/TypeSubstitutionTests.java
+++ b/src/test/java/org/shedlang/compiler/types/TypeSubstitutionTests.java
@@ -36,6 +36,60 @@ public class TypeSubstitutionTests {
     }
 
     @Test
+    public void substitutingFunctionSubstitutesPositionalArgs() {
+        var typeSubstitutionMap = TypeSubstitutionMap.of(
+            Map.entry(TypeParam.forStructType(STRUCT_TYPE, "T"), Types.INT)
+        );
+
+        var result = TypeSubstitution.substitute(
+            FunctionType.builder()
+                .addPositionalParam(TypeParam.forStructType(STRUCT_TYPE, "T"))
+                .build(),
+            typeSubstitutionMap
+        );
+
+        assertThat(result, equalTo(FunctionType.builder()
+            .addPositionalParam(Types.INT)
+            .build()));
+    }
+
+    @Test
+    public void substitutingFunctionSubstitutesNamedArgs() {
+        var typeSubstitutionMap = TypeSubstitutionMap.of(
+            Map.entry(TypeParam.forStructType(STRUCT_TYPE, "T"), Types.INT)
+        );
+
+        var result = TypeSubstitution.substitute(
+            FunctionType.builder()
+                .addNamedParam("x", TypeParam.forStructType(STRUCT_TYPE, "T"))
+                .build(),
+            typeSubstitutionMap
+        );
+
+        assertThat(result, equalTo(FunctionType.builder()
+            .addNamedParam("x", Types.INT)
+            .build()));
+    }
+
+    @Test
+    public void substitutingFunctionSubstitutesReturnType() {
+        var typeSubstitutionMap = TypeSubstitutionMap.of(
+            Map.entry(TypeParam.forStructType(STRUCT_TYPE, "T"), Types.INT)
+        );
+
+        var result = TypeSubstitution.substitute(
+            FunctionType.builder()
+                .withReturnType(TypeParam.forStructType(STRUCT_TYPE, "T"))
+                .build(),
+            typeSubstitutionMap
+        );
+
+        assertThat(result, equalTo(FunctionType.builder()
+            .withReturnType(Types.INT)
+            .build()));
+    }
+
+    @Test
     public void substitutingIntDoesNothing() {
         var typeSubstitutionMap = TypeSubstitutionMap.of(
             Map.entry(TypeParam.forStructType(STRUCT_TYPE, "T"), Types.BOOL)