From 4d131f97a5683e4089cb1c47a302acc0c7ad2bf0 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 19 Feb 2025 19:55:11 +0800 Subject: [PATCH] tests: rm dup code --- src/flow/src/recording_rules.rs | 54 ++++++++++++++------------------- 1 file changed, 23 insertions(+), 31 deletions(-) diff --git a/src/flow/src/recording_rules.rs b/src/flow/src/recording_rules.rs index 24fe0459556e..92c372cc7570 100644 --- a/src/flow/src/recording_rules.rs +++ b/src/flow/src/recording_rules.rs @@ -454,42 +454,34 @@ mod test { #[tokio::test] async fn test_add_filter() { + let testcases = vec![ + ( + "SELECT number FROM numbers_with_ts GROUP BY number","SELECT numbers_with_ts.number FROM numbers_with_ts WHERE (number > 4) GROUP BY numbers_with_ts.number" + ), + ( + "SELECT number FROM numbers_with_ts WHERE number < 2 OR number >10", + "SELECT numbers_with_ts.number FROM numbers_with_ts WHERE (((numbers_with_ts.number < 2) OR (numbers_with_ts.number > 10)) AND (number > 4))" + ), + ( + "SELECT date_bin('5 minutes', ts) as time_window FROM numbers_with_ts GROUP BY time_window", + "SELECT date_bin('5 minutes', numbers_with_ts.ts) AS time_window FROM numbers_with_ts WHERE (number > 4) GROUP BY date_bin('5 minutes', numbers_with_ts.ts)" + ) + ]; use datafusion_expr::{col, lit}; let query_engine = create_test_query_engine(); let ctx = QueryContext::arc(); - let sql = "SELECT number FROM numbers_with_ts GROUP BY number"; - let plan = sql_to_df_plan(ctx.clone(), query_engine.clone(), sql, false) - .await - .unwrap(); - - let mut add_filter = AddFilterRewriter::new(col("number").gt(lit(4u32))); - let plan = plan.rewrite(&mut add_filter).unwrap().data; - let new_sql = df_plan_to_sql(&plan).unwrap(); - assert_eq!( - "SELECT numbers_with_ts.number FROM numbers_with_ts WHERE (number > 4) GROUP BY numbers_with_ts.number", - new_sql - ); - - let sql = "SELECT number FROM numbers_with_ts WHERE number < 2 OR number >10"; - let plan = sql_to_df_plan(ctx.clone(), query_engine.clone(), sql, false) - .await - .unwrap(); - - let mut add_filter = AddFilterRewriter::new(col("number").gt(lit(4u32))); - let plan = plan.rewrite(&mut add_filter).unwrap().data; - let new_sql = df_plan_to_sql(&plan).unwrap(); - assert_eq!("SELECT numbers_with_ts.number FROM numbers_with_ts WHERE (((numbers_with_ts.number < 2) OR (numbers_with_ts.number > 10)) AND (number > 4))", new_sql); - - let sql = "SELECT date_bin('5 minutes', ts) as time_window FROM numbers_with_ts GROUP BY time_window"; - let plan = sql_to_df_plan(ctx.clone(), query_engine.clone(), sql, false) - .await - .unwrap(); + for (before, after) in testcases { + let sql = before; + let plan = sql_to_df_plan(ctx.clone(), query_engine.clone(), sql, false) + .await + .unwrap(); - let mut add_filter = AddFilterRewriter::new(col("number").gt(lit(4u32))); - let plan = plan.rewrite(&mut add_filter).unwrap().data; - let new_sql = df_plan_to_sql(&plan).unwrap(); - assert_eq!("SELECT date_bin('5 minutes', numbers_with_ts.ts) AS time_window FROM numbers_with_ts WHERE (number > 4) GROUP BY date_bin('5 minutes', numbers_with_ts.ts)", new_sql); + let mut add_filter = AddFilterRewriter::new(col("number").gt(lit(4u32))); + let plan = plan.rewrite(&mut add_filter).unwrap().data; + let new_sql = df_plan_to_sql(&plan).unwrap(); + assert_eq!(after, new_sql); + } } #[tokio::test]