diff --git a/pilopt/tests/optimizer.rs b/pilopt/tests/optimizer.rs
index 232bb72817..f1e921b569 100644
--- a/pilopt/tests/optimizer.rs
+++ b/pilopt/tests/optimizer.rs
@@ -282,29 +282,26 @@ fn enum_ref_by_trait() {
     assert_eq!(optimized, expectation);
 }
 
-
 #[test]
 fn handle_array_references_in_prover_functions() {
     let input = r#"namespace N(8);
-        col witness x;
-        col fixed cnt(i) { inc(i) };
-        let inc = |x| x + 1;
-        // these are removed
-        col witness k;
-        col k2 = k;
-        let rec: -> int = || rec();
-        let a: int -> int = |i| b(i + 1);
-        let b: int -> int = |j| 8;
-        // identity
-        [ x ] in [ cnt ];
+        col witness x[1];
+        
+        // non-trivial constraint so that `x[0]` does not get removed.
+        x[0]' = x[0] + 1;
 
+        query |i| {
+            // No-op, but references `x[0]`.
+            let _ = x[0] + 1;
+        };
     "#;
-    let expectation = r#"namespace N(65536);
-    col witness x;
-    col fixed cnt(i) { N::inc(i) };
-    let inc: int -> int = |x| x + 1_int;
-    [N::x] in [N::cnt];
+    let expectation = r#"namespace N(8);
+    col witness x[1];
+    N::x[0]' = N::x[0] + 1;
+    query |i| {
+        let _: expr = N::x[0_int] + 1_expr;
+    };
 "#;
     let optimized = optimize(analyze_string::<GoldilocksField>(input).unwrap()).to_string();
     assert_eq!(optimized, expectation);
-}
\ No newline at end of file
+}