Abstract out making individual reports

keep-starknet-strange · Feb 28, 2024 · 8ad31dc · 8ad31dc
1 parent c0b1f15
commit 8ad31dc
Showing 1 changed file with 40 additions and 36 deletions.
diff --git a/src/actions/mod.rs b/src/actions/mod.rs
@@ -1,4 +1,6 @@
-use starknet::providers::Provider;
+use ::goose::metrics::GooseMetrics;
+use futures::Future;
+use starknet::providers::{jsonrpc::HttpTransport, JsonRpcClient, Provider};
 
 use crate::{
     config::GatlingConfig,
@@ -30,48 +32,27 @@ pub async fn shoot(config: GatlingConfig) -> color_eyre::Result<()> {
     let start_block = shooter.rpc_client().block_number().await?;
 
     if run_erc20 {
-        let start_block = shooter.rpc_client().block_number().await?;
-        let goose_metrics = goose::erc20(&shooter).await?;
-        let end_block = shooter.rpc_client().block_number().await?;
-
-        let mut report =
-            BenchmarkReport::new("Erc20 Transfers", goose_metrics.scenarios[0].counter);
-
-        report
-            .with_block_range(shooter.rpc_client(), start_block + 1, end_block)
-            .await?;
-
-        if config.report.num_blocks != 0 {
-            report
-                .with_last_x_blocks(shooter.rpc_client(), config.report.num_blocks)
-                .await?;
-        }
-
-        report.with_goose_metrics(&goose_metrics)?;
+        let report = make_report_over_bench(
+            goose::erc20(&shooter),
+            "Erc20 Transfers",
+            shooter.rpc_client(),
+            config.report.num_blocks,
+        )
+        .await?;
 
         whole_report.benches.push(report);
     } else {
         log::info!("Skipping erc20 transfers")
     }
 
     if run_erc721 {
-        let start_block = shooter.rpc_client().block_number().await?;
-        let goose_metrics = goose::erc721(&shooter).await?;
-        let end_block = shooter.rpc_client().block_number().await?;
-
-        let mut report = BenchmarkReport::new("Erc721 Mints", goose_metrics.scenarios[0].counter);
-
-        report
-            .with_block_range(shooter.rpc_client(), start_block + 1, end_block)
-            .await?;
-
-        if config.report.num_blocks != 0 {
-            report
-                .with_last_x_blocks(shooter.rpc_client(), config.report.num_blocks)
-                .await?;
-        }
-
-        report.with_goose_metrics(&goose_metrics)?;
+        let report = make_report_over_bench(
+            goose::erc721(&shooter),
+            "Erc721 Mints",
+            shooter.rpc_client(),
+            config.report.num_blocks,
+        )
+        .await?;
 
         whole_report.benches.push(report);
     } else {
@@ -92,3 +73,26 @@ pub async fn shoot(config: GatlingConfig) -> color_eyre::Result<()> {
 
     Ok(())
 }
+
+async fn make_report_over_bench(
+    bench: impl Future<Output = color_eyre::Result<GooseMetrics>>,
+    name: &'static str,
+    rpc_client: &JsonRpcClient<HttpTransport>,
+    num_blocks: u64,
+) -> color_eyre::Result<BenchmarkReport> {
+    let start_block = rpc_client.block_number().await?;
+    let goose_metrics = bench.await?;
+    let end_block = rpc_client.block_number().await?;
+
+    let mut report = BenchmarkReport::new(name, goose_metrics.scenarios[0].counter);
+    report
+        .with_block_range(rpc_client, start_block + 1, end_block)
+        .await?;
+
+    if num_blocks != 0 {
+        report.with_last_x_blocks(rpc_client, num_blocks).await?;
+    }
+
+    report.with_goose_metrics(&goose_metrics)?;
+    Ok(report)
+}