diff --git a/src/circuit_tests/poseidon-digest-test.circom b/src/circuit_tests/poseidon-digest-test.circom
index f8d244f..8bb957c 100644
--- a/src/circuit_tests/poseidon-digest-test.circom
+++ b/src/circuit_tests/poseidon-digest-test.circom
@@ -2,4 +2,19 @@ pragma circom 2.1.0;
 
 include "../../circuits/poseidon-digest.circom";
 
-component main = PoseidonDigest(256, 16);
+template PoseidonDigestTest(BLOCK_SIZE, CHUNK_SIZE) {
+    signal input block[BLOCK_SIZE];
+    signal input hash;
+    signal output hash2;
+
+    component digest = PoseidonDigest(BLOCK_SIZE, CHUNK_SIZE);
+    for (var i = 0; i < BLOCK_SIZE; i++) {
+        digest.block[i] <== block[i];
+    }
+
+    digest.hash === hash; // verify that the hash is correct
+
+    hash2 <== digest.hash;
+}
+
+component main { public [hash] } = PoseidonDigestTest(256, 16);
diff --git a/src/circuit_tests/poseidon-hash-test.circom b/src/circuit_tests/poseidon-hash-test.circom
index 5450e49..57a11e1 100644
--- a/src/circuit_tests/poseidon-hash-test.circom
+++ b/src/circuit_tests/poseidon-hash-test.circom
@@ -12,4 +12,4 @@ template PoseidonHash(SIZE) {
     hasher.out === hash;
 }
 
-component main = PoseidonHash(1);
+component main { public [hash] } = PoseidonHash(1);
diff --git a/src/circuit_tests/utils.rs b/src/circuit_tests/utils.rs
index 8c372c7..a8faf69 100644
--- a/src/circuit_tests/utils.rs
+++ b/src/circuit_tests/utils.rs
@@ -6,16 +6,14 @@ pub fn digest(input: &[U256], chunk_size: Option<usize>) -> U256 {
     let chunks = ((input.len() as f32) / (chunk_size as f32)).ceil() as usize;
     let mut concat: Vec<U256> = vec![];
 
-    let mut i: usize = 0;
-    while i < chunks {
+    for i in 0..chunks {
         let range = (i * chunk_size)..std::cmp::min((i + 1) * chunk_size, input.len());
-        let mut chunk: Vec<U256> = input[range].to_vec();
+        let mut chunk = input[range].to_vec();
         if chunk.len() < chunk_size {
             chunk.resize(chunk_size as usize, uint!(0_U256));
         }
 
         concat.push(hash(chunk.as_slice()));
-        i += chunk_size;
     }
 
     if concat.len() > 1 {