From 835702fdc3fc6ff1fbdc2e2651960ed0990e297a Mon Sep 17 00:00:00 2001
From: Matthew Turner <matthew.m.turner@outlook.com>
Date: Tue, 15 Oct 2024 10:17:20 -0400
Subject: [PATCH] Fixes

---
 src/execution/flightsql.rs            |  2 +-
 src/execution/flightsql_benchmarks.rs |  2 +-
 tests/extension_cases/flightsql.rs    | 42 +++++++++++++++++++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/src/execution/flightsql.rs b/src/execution/flightsql.rs
index a89f5ae..e4d2628 100644
--- a/src/execution/flightsql.rs
+++ b/src/execution/flightsql.rs
@@ -26,7 +26,7 @@ use tonic::{transport::Channel, IntoRequest};
 
 use crate::config::FlightSQLConfig;
 
-use super::local_benchmarks::FlightSQLBenchmarkStats;
+use crate::execution::flightsql_benchmarks::FlightSQLBenchmarkStats;
 
 pub type FlightSQLClient = Mutex<Option<FlightSqlServiceClient<Channel>>>;
 
diff --git a/src/execution/flightsql_benchmarks.rs b/src/execution/flightsql_benchmarks.rs
index 2051af5..392a843 100644
--- a/src/execution/flightsql_benchmarks.rs
+++ b/src/execution/flightsql_benchmarks.rs
@@ -62,7 +62,7 @@ impl FlightSQLBenchmarkStats {
             .map(|d| d.as_nanos())
             .sum::<u128>();
         let percent_of_total = (this_total as f64 / total_duration as f64) * 100.0;
-        LocalDurationsSummary {
+        DurationsSummary {
             min,
             max,
             mean,
diff --git a/tests/extension_cases/flightsql.rs b/tests/extension_cases/flightsql.rs
index 1ac62a2..76f09fd 100644
--- a/tests/extension_cases/flightsql.rs
+++ b/tests/extension_cases/flightsql.rs
@@ -204,3 +204,45 @@ pub async fn test_time_files() {
     assert.stdout(contains_str(expected));
     fixture.shutdown_and_wait().await;
 }
+
+#[test]
+fn test_bench_command() {
+    let assert = Command::cargo_bin("dft")
+        .unwrap()
+        .arg("-c")
+        .arg("SELECT 1")
+        .arg("--bench")
+        .arg("--flightsql")
+        .assert()
+        .success();
+
+    let expected = r##"
+----------------------------
+Benchmark Stats (10 runs)
+----------------------------
+SELECT 1
+----------------------------"##;
+    assert.stdout(contains_str(expected));
+}
+
+#[test]
+fn test_bench_files() {
+    let file = sql_in_file(r#"SELECT 1 + 1;"#);
+
+    let assert = Command::cargo_bin("dft")
+        .unwrap()
+        .arg("-f")
+        .arg(file.path())
+        .arg("--bench")
+        .arg("--flightsql")
+        .assert()
+        .success();
+
+    let expected_err = r##"
+----------------------------
+Benchmark Stats (10 runs)
+----------------------------
+SELECT 1 + 1;
+----------------------------"##;
+    assert.code(0).stdout(contains_str(expected_err));
+}