From 6b662f903bbb09c153a7b5a99b50d77a9d4afe00 Mon Sep 17 00:00:00 2001
From: Noah Citron <noah@jeff.org>
Date: Tue, 31 Jan 2023 18:23:55 -0500
Subject: [PATCH] fix: benches (#184)

---
 benches/file_db.rs | 20 ++++++++++++++------
 benches/harness.rs | 14 +++++++-------
 2 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/benches/file_db.rs b/benches/file_db.rs
index 6d0baec..6f5e65a 100644
--- a/benches/file_db.rs
+++ b/benches/file_db.rs
@@ -1,4 +1,5 @@
 use client::database::Database;
+use config::Config;
 use criterion::{criterion_group, criterion_main, Criterion};
 use helios::prelude::FileDB;
 use tempfile::tempdir;
@@ -17,8 +18,12 @@ pub fn save_checkpoint(c: &mut Criterion) {
     c.bench_function("save_checkpoint", |b| {
         let checkpoint = vec![1, 2, 3];
         b.iter(|| {
-            let temp_dir = tempdir().unwrap().into_path();
-            let db = FileDB::new(temp_dir);
+            let data_dir = Some(tempdir().unwrap().into_path());
+            let config = Config {
+                data_dir,
+                ..Default::default()
+            };
+            let db = FileDB::new(&config).unwrap();
             db.save_checkpoint(checkpoint.clone()).unwrap();
         })
     });
@@ -28,14 +33,17 @@ pub fn save_checkpoint(c: &mut Criterion) {
 pub fn load_checkpoint(c: &mut Criterion) {
     c.bench_function("load_checkpoint", |b| {
         // First write to the db
-        let temp_dir = tempdir().unwrap().into_path();
-        let db = FileDB::new(temp_dir.clone());
-        let written_checkpoint = vec![1, 2, 3];
+        let data_dir = Some(tempdir().unwrap().into_path());
+        let config = Config {
+            data_dir,
+            ..Default::default()
+        };
+        let db = FileDB::new(&config).unwrap();
+        let written_checkpoint = vec![1; 32];
         db.save_checkpoint(written_checkpoint.clone()).unwrap();
 
         // Then read from the db
         b.iter(|| {
-            let db = FileDB::new(temp_dir.clone());
             let checkpoint = db.load_checkpoint().unwrap();
             assert_eq!(checkpoint, written_checkpoint.clone());
         })
diff --git a/benches/harness.rs b/benches/harness.rs
index 1f9b069..4fb463b 100644
--- a/benches/harness.rs
+++ b/benches/harness.rs
@@ -2,12 +2,12 @@
 
 use std::{str::FromStr, sync::Arc};
 
-use ::client::Client;
+use ::client::{database::ConfigDB, Client};
 use ethers::{
     abi::Address,
     types::{H256, U256},
 };
-use helios::{client, config::networks, prelude::FileDB, types::BlockTag};
+use helios::{client, config::networks, types::BlockTag};
 
 /// Fetches the latest mainnet checkpoint from the fallback service.
 ///
@@ -29,11 +29,11 @@ pub async fn fetch_mainnet_checkpoint() -> eyre::Result<H256> {
 /// The client will use `https://www.lightclientdata.org` as the consensus RPC.
 pub fn construct_mainnet_client(
     rt: &tokio::runtime::Runtime,
-) -> eyre::Result<client::Client<client::FileDB>> {
+) -> eyre::Result<client::Client<ConfigDB>> {
     rt.block_on(inner_construct_mainnet_client())
 }
 
-pub async fn inner_construct_mainnet_client() -> eyre::Result<client::Client<client::FileDB>> {
+pub async fn inner_construct_mainnet_client() -> eyre::Result<client::Client<ConfigDB>> {
     let benchmark_rpc_url = std::env::var("MAINNET_RPC_URL")?;
     let mut client = client::ClientBuilder::new()
         .network(networks::Network::MAINNET)
@@ -47,7 +47,7 @@ pub async fn inner_construct_mainnet_client() -> eyre::Result<client::Client<cli
 
 pub async fn construct_mainnet_client_with_checkpoint(
     checkpoint: &str,
-) -> eyre::Result<client::Client<client::FileDB>> {
+) -> eyre::Result<client::Client<ConfigDB>> {
     let benchmark_rpc_url = std::env::var("MAINNET_RPC_URL")?;
     let mut client = client::ClientBuilder::new()
         .network(networks::Network::MAINNET)
@@ -79,7 +79,7 @@ pub fn construct_runtime() -> tokio::runtime::Runtime {
 /// The client will use `http://testing.prater.beacon-api.nimbus.team` as the consensus RPC.
 pub fn construct_goerli_client(
     rt: &tokio::runtime::Runtime,
-) -> eyre::Result<client::Client<client::FileDB>> {
+) -> eyre::Result<client::Client<ConfigDB>> {
     rt.block_on(async {
         let benchmark_rpc_url = std::env::var("GOERLI_RPC_URL")?;
         let mut client = client::ClientBuilder::new()
@@ -96,7 +96,7 @@ pub fn construct_goerli_client(
 /// Gets the balance of the given address on mainnet.
 pub fn get_balance(
     rt: &tokio::runtime::Runtime,
-    client: Arc<Client<FileDB>>,
+    client: Arc<Client<ConfigDB>>,
     address: &str,
 ) -> eyre::Result<U256> {
     rt.block_on(async {