From 0db9021da08d776919a5781d4748f8a9710898cf Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Thu, 2 Sep 2021 16:30:57 +0800
Subject: [PATCH] Update tests

---
 src/executor/select.rs    |  2 +-
 src/query/helper.rs       | 29 +++++++++++++++++-
 tests/relational_tests.rs | 63 +++++++++++++++++++++++++++++++--------
 3 files changed, 79 insertions(+), 15 deletions(-)

diff --git a/src/executor/select.rs b/src/executor/select.rs
index 6196ec73..836ea495 100644
--- a/src/executor/select.rs
+++ b/src/executor/select.rs
@@ -128,7 +128,7 @@ where
     E: EntityTrait,
     F: EntityTrait,
 {
-    fn into_model<M, N>(self) -> Selector<SelectTwoModel<M, N>>
+    pub fn into_model<M, N>(self) -> Selector<SelectTwoModel<M, N>>
     where
         M: FromQueryResult,
         N: FromQueryResult,
diff --git a/src/query/helper.rs b/src/query/helper.rs
index 6ade581a..4ca44905 100644
--- a/src/query/helper.rs
+++ b/src/query/helper.rs
@@ -3,7 +3,7 @@ use crate::{
     RelationDef,
 };
 use sea_query::{
-    Alias, Expr, IntoCondition, SeaRc, SelectExpr, SelectStatement, SimpleExpr, TableRef,
+    Alias, Expr, Iden, IntoCondition, SeaRc, SelectExpr, SelectStatement, SimpleExpr, TableRef,
 };
 pub use sea_query::{Condition, ConditionalStatement, DynIden, JoinType, Order, OrderedStatement};
 
@@ -66,6 +66,33 @@ pub trait QuerySelect: Sized {
         self
     }
 
+    /// Add a select column with prefixed alias
+    /// ```
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, DbBackend};
+    ///
+    /// assert_eq!(
+    ///     cake::Entity::find()
+    ///         .select_only()
+    ///         .column_as_prefixed(cake::Column::Id.count(), "A_", cake::Column::Id)
+    ///         .build(DbBackend::Postgres)
+    ///         .to_string(),
+    ///     r#"SELECT COUNT("cake"."id") AS "A_id" FROM "cake""#
+    /// );
+    /// ```
+    fn column_as_prefixed<C, I>(mut self, col: C, prefix: &str, alias: I) -> Self
+    where
+        C: IntoSimpleExpr,
+        I: Iden,
+    {
+        self.query().expr(SelectExpr {
+            expr: col.into_simple_expr(),
+            alias: Some(SeaRc::new(Alias::new(
+                vec![prefix, alias.to_string().as_str()].join("").as_str(),
+            ))),
+        });
+        self
+    }
+
     /// Add a group by column
     /// ```
     /// use sea_orm::{entity::*, query::*, tests_cfg::cake, DbBackend};
diff --git a/tests/relational_tests.rs b/tests/relational_tests.rs
index 1ddee902..f67d4e88 100644
--- a/tests/relational_tests.rs
+++ b/tests/relational_tests.rs
@@ -652,23 +652,60 @@ pub async fn linked() -> Result<(), DbErr> {
     .save(&ctx.db)
     .await?;
 
-    /*
-       SELECT `baker`.`id` AS `A_id`, `baker`.`name` AS `A_name`,
-           `baker`.`contact_details` AS `A_contact_details`,
-           `baker`.`bakery_id` AS `A_bakery_id`, `customer`.`id` AS `B_id`,
-           `customer`.`name` AS `B_name`, `customer`.`notes` AS `B_notes`
-       FROM `baker`
-       LEFT JOIN `cakes_bakers` ON `baker`.`id` = `cakes_bakers`.`baker_id`
-       LEFT JOIN `cake` ON `cakes_bakers`.`cake_id` = `cake`.`id`
-       LEFT JOIN `lineitem` ON `cake`.`id` = `lineitem`.`cake_id`
-       LEFT JOIN `order` ON `lineitem`.`order_id` = `order`.`id`
-       LEFT JOIN `customer` ON `order`.`customer_id` = `customer`.`id`
-    */
+    #[derive(Debug, FromQueryResult, PartialEq)]
+    struct BakerLite {
+        name: String,
+    }
+
+    #[derive(Debug, FromQueryResult, PartialEq)]
+    struct CustomerLite {
+        name: String,
+    }
+
     let baked_for_customers = Baker::find()
         .find_also_linked(baker::BakedForCustomer)
+        .select_only()
+        .column_as_prefixed(baker::Column::Name, "A_", baker::Column::Name)
+        .column_as_prefixed(customer::Column::Name, "B_", customer::Column::Name)
+        .group_by(baker::Column::Id)
+        .group_by(customer::Column::Id)
+        .group_by(baker::Column::Name)
+        .group_by(customer::Column::Name)
+        .order_by_asc(baker::Column::Id)
+        .order_by_asc(customer::Column::Id)
+        .into_model::<BakerLite, CustomerLite>()
         .all(&ctx.db)
         .await?;
-    println!("{:#?}", baked_for_customers);
+
+    assert_eq!(
+        baked_for_customers,
+        vec![
+            (
+                BakerLite {
+                    name: "Baker Bob".to_owned(),
+                },
+                Some(CustomerLite {
+                    name: "Kara".to_owned(),
+                })
+            ),
+            (
+                BakerLite {
+                    name: "Baker Bobby".to_owned(),
+                },
+                Some(CustomerLite {
+                    name: "Kate".to_owned(),
+                })
+            ),
+            (
+                BakerLite {
+                    name: "Baker Bobby".to_owned(),
+                },
+                Some(CustomerLite {
+                    name: "Kara".to_owned(),
+                })
+            ),
+        ]
+    );
 
     ctx.delete().await;