diff --git a/test/test_expect.jl b/test/test_expect.jl
index 75ed8504..c654138e 100644
--- a/test/test_expect.jl
+++ b/test/test_expect.jl
@@ -29,11 +29,13 @@ using Test: @test, @testset
   s = siteinds("S=1/2", g)
   rng = StableRNG(1234)
   ψ = random_tensornetwork(rng, s; link_space=χ)
-  cache_construction_function =
-    f -> BeliefPropagationCache(
-      f; partitioned_vertices=group(v -> (original_state_vertex(f, v)[1]), vertices(f))
-    )
-  sz_bp = expect(ψ, "Sz"; alg="bp", cache_construction_function)
+  quadratic_form_vertices = reduce(
+    vcat, [[(v, "ket"), (v, "bra"), (v, "operator")] for v in vertices(ψ)]
+  )
+  cache_construction_kwargs = (;
+    partitioned_vertices=group(v -> first(first(v)), quadratic_form_vertices)
+  )
+  sz_bp = expect(ψ, "Sz"; alg="bp", cache_construction_kwargs)
   sz_exact = expect(ψ, "Sz"; alg="exact")
   @test sz_bp ≈ sz_exact