Fix some tests

blolt · May 17, 2024 · 82f598e · 82f598e
1 parent 0c59ef2
commit 82f598e
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 12 deletions.
diff --git a/src/graph_engine.jl b/src/graph_engine.jl
@@ -292,7 +292,7 @@ function unroll(proxylabel::ProxyLabel, proxied::ProxyLabel, index, maycreate, l
 end
 
 function unroll(proxylabel::ProxyLabel, something::Any, index, maycreate, liftedindex)
-    return something
+    return checked_getindex(something, index)
 end
 
 checked_getindex(something, index::FunctionalIndex) = Base.getindex(something, index)

diff --git a/test/graph_construction_tests.jl b/test/graph_construction_tests.jl
@@ -121,7 +121,7 @@ end
 @testitem "Simple model with lazy data (vector) creation" begin
     using Distributions
     import GraphPPL:
-        create_model, getorcreate!, LazyIndex, NodeCreationOptions, MissingCollection, index, getproperties, is_kind, VariableRef, datalabel
+        create_model, getorcreate!, LazyIndex, NodeCreationOptions, MissingCollection, index, getproperties, is_kind, VariableRef, datalabel, VariableKindData
 
     include("testutils.jl")
 
@@ -138,9 +138,9 @@ end
 
     @testset for n in 5:10
         model = create_model(simple_model_3(n = n)) do model, ctx
-            T = datalabel(model, ctx, NodeCreationOptions(kind = :data_for_T), :T)
-            y = datalabel(model, ctx, NodeCreationOptions(kind = :data_for_y), :y)
-            Σ = datalabel(model, ctx, NodeCreationOptions(kind = :data_for_Σ), :Σ)
+            T = datalabel(model, ctx, NodeCreationOptions(kind = VariableKindData), :T)
+            y = datalabel(model, ctx, NodeCreationOptions(kind = VariableKindData), :y)
+            Σ = datalabel(model, ctx, NodeCreationOptions(kind = VariableKindData), :Σ)
             return (T = T, y = y, Σ = Σ)
         end
 
@@ -151,9 +151,9 @@ end
         @test length(collect(filter(as_variable(:y), model))) === n
 
         # test that options are preserved
-        @test all(label -> is_kind(getproperties(model[label]), :data_for_T), collect(filter(as_variable(:T), model)))
-        @test all(label -> is_kind(getproperties(model[label]), :data_for_y), collect(filter(as_variable(:y), model)))
-        @test all(label -> is_kind(getproperties(model[label]), :data_for_Σ), collect(filter(as_variable(:Σ), model)))
+        @test all(label -> is_kind(getproperties(model[label]), VariableKindData), collect(filter(as_variable(:T), model)))
+        @test all(label -> is_kind(getproperties(model[label]), VariableKindData), collect(filter(as_variable(:y), model)))
+        @test all(label -> is_kind(getproperties(model[label]), VariableKindData), collect(filter(as_variable(:Σ), model)))
 
         # test that indices are of expected shape
         Tsindices = map((label) -> index(getproperties(model[label])), collect(filter(as_variable(:T), model)))

diff --git a/test/graph_engine_tests.jl b/test/graph_engine_tests.jl
@@ -632,7 +632,7 @@ end
 end
 
 @testitem "Lift index" begin
-    import GraphPPL: lift_index, True, False
+    import GraphPPL: lift_index, True, False, checked_getindex
 
     @test lift_index(True(), nothing, nothing) === nothing
     @test lift_index(True(), (1,), nothing) === (1,)
@@ -653,11 +653,10 @@ end
     GraphPPL.is_proxied(::Type{LiftingTest}) = GraphPPL.True()
 
     function GraphPPL.unroll(proxy::ProxyLabel, ::LiftingTest, index, maycreate, liftedindex)
-        @test maycreate === True()
         if liftedindex === nothing
-            return "Hello"
+            return checked_getindex("Hello", index)
         else
-            return "World"
+            return checked_getindex("World", index)
         end
     end