Add type info for sum types

1 files changed, 12 insertions, 3 deletions

diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerSumTypeDefinitionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerSumTypeDefinitionTests.java
index 0a551c95..d70d3660 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerSumTypeDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerSumTypeDefinitionTests.java
@@ -15,8 +15,7 @@ import java.util.Map;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.isTypedSumTypeDefinitionNode;
 import static org.zwobble.precisely.AssertThat.assertThat;
-import static org.zwobble.precisely.Matchers.equalTo;
-import static org.zwobble.precisely.Matchers.isOptionalOf;
+import static org.zwobble.precisely.Matchers.*;
 
 public class TypeCheckerSumTypeDefinitionTests {
     private static final NamespaceName NAMESPACE_NAME = NamespaceName.of("Example");
@@ -57,7 +56,7 @@ public class TypeCheckerSumTypeDefinitionTests {
         var untyped = UntypedSumTypeDefinitionNode.builder()
             .withName("ExampleSumType")
             .addVariant(Untyped.typeLevelReference("A"))
-            .addVariant(Untyped.typeLevelReference("A"))
+            .addVariant(Untyped.typeLevelReference("B"))
             .build();
         var variantType1 = Types.struct(NAMESPACE_NAME, "A");
         var variantType2 = Types.struct(NAMESPACE_NAME, "B");
@@ -72,6 +71,16 @@ public class TypeCheckerSumTypeDefinitionTests {
         assertThat(typed.value(), isTypedSumTypeDefinitionNode(
             Types.sum(NAMESPACE_NAME, "ExampleSumType")
         ));
+        var sumTypeInfo = typed.context().toTypeInfos()
+            .sumTypeInfo(Types.sum(NAMESPACE_NAME, "ExampleSumType"))
+            .orElseThrow();
+        assertThat(
+            sumTypeInfo,
+            has("variants", x -> x.variants(), isSequence(
+                equalTo(variantType1),
+                equalTo(variantType2)
+            ))
+        );
     }
 
     private static TypeCheckerResult<? extends TypedNamespaceStatementNode> typeCheckNamespaceStatement(