Collect mapping of variants to sum types

1 files changed, 12 insertions, 6 deletions

diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java
index 1d968d6..fdf88ee 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerSumDefinitionTests.java
@@ -49,13 +49,15 @@ public class TypeCheckerSumDefinitionTests {
             DocComment.EMPTY,
             NullSource.INSTANCE
         );
-        var context = TypeCheckerContextArb.namespaceContext(NamespaceName.of("a", "b"));
-        context.declare(rectangleType.name());
-        context.define(rectangleType.name(), new TypeLevelValueType(rectangleType));
-        context.declare(triangleType.name());
-        context.define(triangleType.name(), new TypeLevelValueType(triangleType));
+        var globalContext = TypeCheckerGlobalContext.initial();
+        var namespaceContext = globalContext
+            .enterNamespace(NamespaceName.of("a", "b"));
+        namespaceContext.declare(rectangleType.name());
+        namespaceContext.define(rectangleType.name(), new TypeLevelValueType(rectangleType));
+        namespaceContext.declare(triangleType.name());
+        namespaceContext.define(triangleType.name(), new TypeLevelValueType(triangleType));
 
-        var typed = typeCheckNamespaceStatement(untyped, context);
+        var typed = typeCheckNamespaceStatement(untyped, namespaceContext);
 
         assertThat(typed, instanceOf(
             TypedSumDefinitionNode.class,
@@ -76,5 +78,9 @@ public class TypeCheckerSumDefinitionTests {
                 )
             )
         ));
+        var sumType = new SumType(NamespaceName.of("a", "b"), "X");
+        var typesInfo = globalContext.toTypesInfo();
+        assertThat(typesInfo.variantOf(rectangleType), isSequence(equalTo(sumType)));
+        assertThat(typesInfo.variantOf(triangleType), isSequence(equalTo(sumType)));
     }
 }