chore(coprocessor): update unit-tests in zkproof to use PostgresPoolManager

Author: goshawk-3

coprocessor/fhevm-engine/zkproof-worker/src/tests/mod.rs

@@ -8,7 +8,8 @@ mod utils;
 #[tokio::test]
 #[serial(db)]
 async fn test_verify_proof() {
-    let (pool, _instance) = utils::setup().await.expect("valid setup");
+    let (pool_mngr, _instance) = utils::setup().await.expect("valid setup");
+    let pool = pool_mngr.pool();
 
     // Generate Valid ZkPok
     let aux: (crate::auxiliary::ZkData, [u8; 92]) =
@@ -42,7 +43,8 @@ async fn test_verify_proof() {
 #[tokio::test]
 #[serial(db)]
 async fn test_verify_empty_input_list() {
-    let (pool, _instance) = utils::setup().await.expect("valid setup");
+    let (pool_mngr, _instance) = utils::setup().await.expect("valid setup");
+    let pool = pool_mngr.pool();
 
     let aux: (crate::auxiliary::ZkData, [u8; 92]) =
         utils::aux_fixture(ACL_CONTRACT_ADDR.to_owned());
@@ -61,7 +63,8 @@ async fn test_verify_empty_input_list() {
 #[tokio::test]
 #[serial(db)]
 async fn test_max_input_index() {
-    let (db, _instance) = utils::setup().await.expect("valid setup");
+    let (pool_mngr, _instance) = utils::setup().await.expect("valid setup");
+    let pool = pool_mngr.pool();
 
     let aux: (crate::auxiliary::ZkData, [u8; 92]) =
         utils::aux_fixture(ACL_CONTRACT_ADDR.to_owned());
@@ -70,11 +73,11 @@ async fn test_max_input_index() {
     let inputs = vec![utils::ZkInput::U8(1); MAX_INPUT_INDEX as usize + 2];
 
     assert!(!utils::is_valid(
-        &db,
+        &pool,
         utils::insert_proof(
-            &db,
+            &pool,
             101,
-            &utils::generate_zk_pok_with_inputs(&db, &aux.1, &inputs).await,
+            &utils::generate_zk_pok_with_inputs(&pool, &aux.1, &inputs).await,
             &aux.0
         )
         .await
@@ -87,11 +90,11 @@ async fn test_max_input_index() {
     // Test with highest number of inputs - 255
     let inputs = vec![utils::ZkInput::U64(2); MAX_INPUT_INDEX as usize + 1];
     assert!(utils::is_valid(
-        &db,
+        &pool,
         utils::insert_proof(
-            &db,
+            &pool,
             102,
-            &utils::generate_zk_pok_with_inputs(&db, &aux.1, &inputs).await,
+            &utils::generate_zk_pok_with_inputs(&pool, &aux.1, &inputs).await,
             &aux.0
         )
         .await

coprocessor/fhevm-engine/zkproof-worker/src/tests/utils.rs

@@ -1,3 +1,4 @@
+use fhevm_engine_common::pg_pool::PostgresPoolManager;
 use fhevm_engine_common::{tenant_keys, utils::safe_serialize};
 use std::sync::Arc;
 use std::time::{Duration, SystemTime};
@@ -7,7 +8,7 @@ use tokio::time::sleep;
 
 use crate::auxiliary::ZkData;
 
-pub async fn setup() -> anyhow::Result<(sqlx::PgPool, DBInstance)> {
+pub async fn setup() -> anyhow::Result<(PostgresPoolManager, DBInstance)> {
     let _ = tracing_subscriber::fmt().json().with_level(true).try_init();
     let test_instance = test_harness::instance::setup_test_db(ImportMode::WithKeysNoSns)
         .await
@@ -22,27 +23,35 @@ pub async fn setup() -> anyhow::Result<(sqlx::PgPool, DBInstance)> {
         worker_thread_count: 1,
     };
 
-    let pool = sqlx::postgres::PgPoolOptions::new()
-        .max_connections(10)
-        .connect(&conf.database_url)
-        .await?;
+    let pool_mngr = PostgresPoolManager::connect_pool(
+        test_instance.parent_token.child_token(),
+        conf.database_url.as_str(),
+        Duration::from_secs(15),
+        10,
+        Duration::from_secs(2),
+        None, // TODO: conf.pg_auto_explain_with_min_duration
+    )
+    .await
+    .unwrap();
+
+    let pmngr = pool_mngr.clone();
 
     sqlx::query("TRUNCATE TABLE verify_proofs")
-        .execute(&pool)
+        .execute(&pmngr.pool())
         .await
         .unwrap();
 
     let last_active_at = Arc::new(RwLock::new(SystemTime::now()));
-    let db_pool = pool.clone();
+
     tokio::spawn(async move {
-        crate::verifier::execute_verify_proofs_loop(db_pool, conf.clone(), last_active_at.clone())
+        crate::verifier::execute_verify_proofs_loop(pmngr, conf.clone(), last_active_at.clone())
             .await
             .unwrap();
     });
 
     sleep(Duration::from_secs(2)).await;
 
-    Ok((pool, test_instance))
+    Ok((pool_mngr, test_instance))
 }
 
 /// Checks if the proof is valid by querying the database continuously.

coprocessor/fhevm-engine/zkproof-worker/src/verifier.rs

@@ -30,7 +30,7 @@ use fhevm_engine_common::healthz_server::{HealthCheckService, HealthStatus, Vers
 use tokio::time::interval;
 use tokio::{select, time::Duration};
 use tokio_util::sync::CancellationToken;
-use tracing::{error, info};
+use tracing::{debug, error, info};
 
 const MAX_CACHED_TENANT_KEYS: usize = 100;
 const EVENT_CIPHERTEXT_COMPUTED: &str = "event_ciphertext_computed";
@@ -205,7 +205,7 @@ async fn execute_worker(
                 };
             },
             _ = idle_event.tick() => {
-                info!("Polling timeout, rechecking for requests");
+                debug!("Polling timeout, rechecking for requests");
             },
             _ = token.cancelled() => {
                 info!("Cancellation requested, stopping worker");