From 897ed7ea5b84dc568c7d4291c25b4c50aae37cd2 Mon Sep 17 00:00:00 2001
From: Denis Smirnov <sd@picodata.io>
Date: Fri, 12 Aug 2022 09:02:53 +0700
Subject: [PATCH] refactoring(test): move bucket_id tests to a separate folder

---
 src/executor/tests.rs           | 30 ++-------------
 src/executor/tests/bucket_id.rs | 67 +++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 27 deletions(-)
 create mode 100644 src/executor/tests/bucket_id.rs

diff --git a/src/executor/tests.rs b/src/executor/tests.rs
index ce0e6c9939..1a029cd127 100644
--- a/src/executor/tests.rs
+++ b/src/executor/tests.rs
@@ -766,33 +766,6 @@ fn sharding_column2_test() {
     assert_eq!(expected, result);
 }
 
-#[test]
-fn bucket1_test() {
-    let sql = r#"SELECT *, "bucket_id" FROM "t1""#;
-    let coordinator = RouterRuntimeMock::new();
-
-    let mut query = Query::new(&coordinator, sql, &[]).unwrap();
-    let result = *query
-        .dispatch()
-        .unwrap()
-        .downcast::<ProducerResult>()
-        .unwrap();
-
-    let mut expected = ProducerResult::new();
-
-    expected.rows.push(vec![
-        Value::String(format!("Execute query on all buckets")),
-        Value::String(String::from(PatternWithParams::new(
-            format!(
-                "{}",
-                r#"SELECT "t1"."a", "t1"."b", "t1"."bucket_id" FROM "t1""#,
-            ),
-            vec![],
-        ))),
-    ]);
-    assert_eq!(expected, result);
-}
-
 #[test]
 fn insert1_test() {
     let sql = r#"insert into "t" ("b") select "a" from "t"
@@ -1372,5 +1345,8 @@ fn get_motion_policy(plan: &Plan, motion_id: usize) -> &MotionPolicy {
 #[cfg(test)]
 mod between;
 
+#[cfg(test)]
+mod bucket_id;
+
 #[cfg(test)]
 mod not_in;
diff --git a/src/executor/tests/bucket_id.rs b/src/executor/tests/bucket_id.rs
new file mode 100644
index 0000000000..8d04341aa3
--- /dev/null
+++ b/src/executor/tests/bucket_id.rs
@@ -0,0 +1,67 @@
+use pretty_assertions::assert_eq;
+
+use crate::executor::engine::cartridge::backend::sql::ir::PatternWithParams;
+use crate::executor::engine::mock::RouterRuntimeMock;
+use crate::executor::result::ProducerResult;
+use crate::ir::value::Value;
+
+use super::*;
+
+#[test]
+fn bucket1_test() {
+    let sql = r#"SELECT *, "bucket_id" FROM "t1""#;
+    let coordinator = RouterRuntimeMock::new();
+
+    let mut query = Query::new(&coordinator, sql, &[]).unwrap();
+    let result = *query
+        .dispatch()
+        .unwrap()
+        .downcast::<ProducerResult>()
+        .unwrap();
+
+    let mut expected = ProducerResult::new();
+
+    expected.rows.push(vec![
+        Value::String(format!("Execute query on all buckets")),
+        Value::String(String::from(PatternWithParams::new(
+            format!(
+                "{}",
+                r#"SELECT "t1"."a", "t1"."b", "t1"."bucket_id" FROM "t1""#,
+            ),
+            vec![],
+        ))),
+    ]);
+    assert_eq!(expected, result);
+}
+
+#[test]
+fn bucket2_test() {
+    let sql = r#"SELECT "a", "bucket_id", "b" FROM "t1"
+        WHERE "a" = 1 AND "b" = 2"#;
+    let coordinator = RouterRuntimeMock::new();
+
+    let mut query = Query::new(&coordinator, sql, &[]).unwrap();
+    let result = *query
+        .dispatch()
+        .unwrap()
+        .downcast::<ProducerResult>()
+        .unwrap();
+
+    let mut expected = ProducerResult::new();
+    let param1 = Value::from(1_u64);
+    let param2 = Value::from(2_u64);
+    let bucket = query.coordinator.determine_bucket_id(&[&param1, &param2]);
+
+    expected.rows.push(vec![
+        Value::String(format!("Execute query on a bucket [{}]", bucket)),
+        Value::String(String::from(PatternWithParams::new(
+            format!(
+                "{} {}",
+                r#"SELECT "t1"."a", "t1"."bucket_id", "t1"."b" FROM "t1""#,
+                r#"WHERE ("t1"."a", "t1"."b") = (?, ?)"#,
+            ),
+            vec![param1, param2],
+        ))),
+    ]);
+    assert_eq!(expected, result);
+}
-- 
GitLab