optimize code

Lordworms · Jul 12, 2024 · d557814 · d557814
1 parent 5e5b93d
commit d557814
Showing 1 changed file with 16 additions and 50 deletions.
diff --git a/datafusion/substrait/tests/cases/consumer_integration.rs b/datafusion/substrait/tests/cases/consumer_integration.rs
@@ -40,53 +40,6 @@ mod tests {
         }
         Ok(ctx)
     }
-    /// query 12 is missing, will be added later
-    async fn create_context_tpch13() -> Result<SessionContext> {
-        let ctx = SessionContext::new();
-
-        let registrations = vec![
-            ("FILENAME_PLACEHOLDER_0", "tests/testdata/tpch/customer.csv"),
-            ("FILENAME_PLACEHOLDER_1", "tests/testdata/tpch/orders.csv"),
-        ];
-
-        for (table_name, file_path) in registrations {
-            register_csv(&ctx, table_name, file_path).await?;
-        }
-
-        Ok(ctx)
-    }
-
-    async fn create_context_tpch14() -> Result<SessionContext> {
-        let ctx = SessionContext::new();
-
-        let registrations = vec![
-            ("FILENAME_PLACEHOLDER_0", "tests/testdata/tpch/lineitem.csv"),
-            ("FILENAME_PLACEHOLDER_1", "tests/testdata/tpch/part.csv"),
-        ];
-
-        for (table_name, file_path) in registrations {
-            register_csv(&ctx, table_name, file_path).await?;
-        }
-
-        Ok(ctx)
-    }
-
-    async fn create_context_tpch16() -> Result<SessionContext> {
-        let ctx = SessionContext::new();
-
-        let registrations = vec![
-            ("FILENAME_PLACEHOLDER_0", "tests/testdata/tpch/partsupp.csv"),
-            ("FILENAME_PLACEHOLDER_1", "tests/testdata/tpch/part.csv"),
-            ("FILENAME_PLACEHOLDER_2", "tests/testdata/tpch/supplier.csv"),
-        ];
-
-        for (table_name, file_path) in registrations {
-            register_csv(&ctx, table_name, file_path).await?;
-        }
-
-        Ok(ctx)
-    }
-
     #[tokio::test]
     async fn tpch_test_1() -> Result<()> {
         let ctx = create_context(vec![(
@@ -364,7 +317,11 @@ mod tests {
     // missing query 12
     #[tokio::test]
     async fn tpch_test_13() -> Result<()> {
-        let ctx = create_context_tpch13().await?;
+        let ctx = create_context(vec![
+            ("FILENAME_PLACEHOLDER_0", "tests/testdata/tpch/customer.csv"),
+            ("FILENAME_PLACEHOLDER_1", "tests/testdata/tpch/orders.csv"),
+        ])
+        .await?;
         let path = "tests/testdata/tpch_substrait_plans/query_13.json";
         let proto = serde_json::from_reader::<_, Plan>(BufReader::new(
             File::open(path).expect("file not found"),
@@ -388,7 +345,11 @@ mod tests {
 
     #[tokio::test]
     async fn tpch_test_14() -> Result<()> {
-        let ctx = create_context_tpch14().await?;
+        let ctx = create_context(vec![
+            ("FILENAME_PLACEHOLDER_0", "tests/testdata/tpch/lineitem.csv"),
+            ("FILENAME_PLACEHOLDER_1", "tests/testdata/tpch/part.csv"),
+        ])
+        .await?;
         let path = "tests/testdata/tpch_substrait_plans/query_14.json";
         let proto = serde_json::from_reader::<_, Plan>(BufReader::new(
             File::open(path).expect("file not found"),
@@ -409,7 +370,12 @@ mod tests {
     // query 15 is missing
     #[tokio::test]
     async fn tpch_test_16() -> Result<()> {
-        let ctx = create_context_tpch16().await?;
+        let ctx = create_context(vec![
+            ("FILENAME_PLACEHOLDER_0", "tests/testdata/tpch/partsupp.csv"),
+            ("FILENAME_PLACEHOLDER_1", "tests/testdata/tpch/part.csv"),
+            ("FILENAME_PLACEHOLDER_2", "tests/testdata/tpch/supplier.csv"),
+        ])
+        .await?;
         let path = "tests/testdata/tpch_substrait_plans/query_16.json";
         let proto = serde_json::from_reader::<_, Plan>(BufReader::new(
             File::open(path).expect("file not found"),