2 files changed, 8 insertions, 2 deletions
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerNamespaceTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerNamespaceTests.java
index 34179cb2..508bf6cd 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerNamespaceTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerNamespaceTests.java
@@ -486,8 +486,11 @@ public class TypeCheckerNamespaceTests {
 
         var result = TypeChecker.typeCheckNamespace(untyped, context);
 
+        var namespaceTypeInfo = result.context().toTypeInfos()
+            .namespaceTypeInfo(namespaceName)
+            .orElseThrow();
         assertThat(
-            result.context().namespaceTypeInfo().get(namespaceName),
+            namespaceTypeInfo,
             equalTo(NamespaceTypeInfo.fromFields(P.orderedMap(
                 Map.entry("f", new FunctionType(List.of(), List.of(), Effects.NONE, Types.UNIT))
             )))
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java
index 5d53301d..53c5d5e8 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerStructDefinitionTests.java
@@ -114,8 +114,11 @@ public class TypeCheckerStructDefinitionTests {
 
         var typed = typeCheckNamespaceStatement(untyped, context);
 
+        var structTypeInfo = typed.context().toTypeInfos()
+            .structTypeInfo(Types.struct(NAMESPACE_NAME, "ExampleStruct"))
+            .orElseThrow();
         assertThat(
-            typed.context().structTypeInfos().get(Types.struct(NAMESPACE_NAME, "ExampleStruct")),
+            structTypeInfo,
             // TODO: builder for OrderedPMap
             has("fields", x -> x.fields().stream().toList(), isSequence(
                 equalTo(Map.entry("b", Types.BOOL)),