From 527b389e17bd4bfd0d81ba484ecbc9d3bb342382 Mon Sep 17 00:00:00 2001
From: Penelope Yong <penelopeysm@gmail.com>
Date: Thu, 28 Nov 2024 14:56:13 +0000
Subject: [PATCH] Don't test setchildcontext on leaf contexts

---
 src/test_utils/contexts.jl | 29 +++++++++++++++--------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/src/test_utils/contexts.jl b/src/test_utils/contexts.jl
index 4748a221c..4de08bf62 100644
--- a/src/test_utils/contexts.jl
+++ b/src/test_utils/contexts.jl
@@ -71,6 +71,15 @@ function test_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Mod
     node_trait isa Union{DynamicPPL.IsLeaf,DynamicPPL.IsParent} ||
         throw(ValueError("Invalid NodeTrait: $node_trait"))
 
+    # To see change, let's make sure we're using a different leaf context than the current.
+    leafcontext_new = if DynamicPPL.leafcontext(context) isa DefaultContext
+        PriorContext()
+    else
+        DefaultContext()
+    end
+    @test DynamicPPL.leafcontext(DynamicPPL.setleafcontext(context, leafcontext_new)) ==
+        leafcontext_new
+
     # The interface methods.
     if node_trait isa DynamicPPL.IsParent
         # `childcontext` and `setchildcontext`
@@ -79,22 +88,14 @@ function test_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Mod
         @test DynamicPPL.childcontext(
             DynamicPPL.setchildcontext(context, childcontext_new)
         ) == childcontext_new
+        # Setting the child context to a leaf should now change the leafcontext
+        # accordingly.
+        context_with_new_leaf = DynamicPPL.setchildcontext(context, leafcontext_new)
+        @test DynamicPPL.childcontext(context_with_new_leaf) ===
+            DynamicPPL.leafcontext(context_with_new_leaf) ===
+            leafcontext_new
     end
 
-    # To see change, let's make sure we're using a different leaf context than the current.
-    leafcontext_new = if DynamicPPL.leafcontext(context) isa DefaultContext
-        PriorContext()
-    else
-        DefaultContext()
-    end
-    @test DynamicPPL.leafcontext(DynamicPPL.setleafcontext(context, leafcontext_new)) ==
-        leafcontext_new
-
-    # Setting the child context to a leaf should now change the leafcontext accordingly.
-    context_with_new_leaf = DynamicPPL.setchildcontext(context, leafcontext_new)
-    @test DynamicPPL.childcontext(context_with_new_leaf) ===
-        DynamicPPL.leafcontext(context_with_new_leaf) ===
-        leafcontext_new
 
     # Make sure that the we can evaluate the model with the context (i.e. that none of the tilde-functions are incorrectly overloaded).
     # The tilde-pipeline contains two different paths: with `SamplingContext` as a parent, and without it.