diff --git a/backend/zk_canister_v2/src/proof.rs b/backend/zk_canister_v2/src/proof.rs
index 8111cfb..b10158e 100644
--- a/backend/zk_canister_v2/src/proof.rs
+++ b/backend/zk_canister_v2/src/proof.rs
@@ -376,10 +376,16 @@ mod tests {
 
     #[test]
     fn test_public_inputs() {
-        let value = Fr::from(100u64);
-        let inputs = create_public_inputs(value.as_int()).expect("Failed to create public inputs");
+        let input = TokenOwnershipInput {
+            balance: 100,
+            min_range: 0,
+            max_range: 200,
+        };
+
+        let inputs = create_public_inputs(&input).expect("Failed to create public inputs");
+
         assert_eq!(inputs.len(), 3);
-        assert_eq!(inputs[0], value);
+        assert_eq!(inputs[0], Fr::from(input.balance));
     }
 
     #[test]
@@ -417,7 +423,7 @@ mod tests {
     }
 
     #[test]
-    fn test_invalid_proof() {
+    fn test_invalid_proof_mismatched_input() {
         let k = 4;
         let mut rng = StdRng::from_entropy();
         let params = ParamsKZG::<Bn256>::setup(k, &mut rng);