From 868a469de0c89a033da886e043a714444eeabb6e Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Wed, 20 Oct 2021 18:02:23 +0800
Subject: [PATCH] Test integer enum

---
 tests/active_enum_tests.rs           | 36 +++++++++++++++++++++++-----
 tests/common/features/active_enum.rs | 23 ++++++++++++++++--
 tests/common/features/schema.rs      | 13 ++++------
 tests/common/setup/mod.rs            |  2 +-
 4 files changed, 57 insertions(+), 17 deletions(-)

diff --git a/tests/active_enum_tests.rs b/tests/active_enum_tests.rs
index 90472fde4..568524814 100644
--- a/tests/active_enum_tests.rs
+++ b/tests/active_enum_tests.rs
@@ -19,19 +19,43 @@ async fn main() -> Result<(), DbErr> {
 }
 
 pub async fn insert_active_enum(db: &DatabaseConnection) -> Result<(), DbErr> {
-    active_enum::ActiveModel {
-        category: Set(active_enum::Category::Big),
+    use active_enum::*;
+
+    let am = ActiveModel {
+        category: Set(None),
+        color: Set(None),
+        // tea: Set(None),
         ..Default::default()
     }
     .insert(db)
     .await?;
 
     assert_eq!(
-        active_enum::Entity::find().one(db).await?.unwrap(),
-        active_enum::Model {
+        Entity::find().one(db).await?.unwrap(),
+        Model {
+            id: 1,
+            category: None,
+            color: None,
+            // tea: None,
+        }
+    );
+
+    ActiveModel {
+        category: Set(Some(Category::Big)),
+        color: Set(Some(Color::Black)),
+        // tea: Set(Some(Tea::EverydayTea)),
+        ..am
+    }
+    .save(db)
+    .await?;
+
+    assert_eq!(
+        Entity::find().one(db).await?.unwrap(),
+        Model {
             id: 1,
-            category: active_enum::Category::Big,
-            category_opt: None,
+            category: Some(Category::Big),
+            color: Some(Color::Black),
+            // tea: Some(Tea::EverydayTea),
         }
     );
 
diff --git a/tests/common/features/active_enum.rs b/tests/common/features/active_enum.rs
index e1aef0fed..d7b15443e 100644
--- a/tests/common/features/active_enum.rs
+++ b/tests/common/features/active_enum.rs
@@ -5,8 +5,9 @@ use sea_orm::entity::prelude::*;
 pub struct Model {
     #[sea_orm(primary_key)]
     pub id: i32,
-    pub category: Category,
-    pub category_opt: Option<Category>,
+    pub category: Option<Category>,
+    pub color: Option<Color>,
+    // pub tea: Option<Tea>,
 }
 
 #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
@@ -22,3 +23,21 @@ pub enum Category {
     #[sea_orm(string_value = "S")]
     Small,
 }
+
+#[derive(Debug, Clone, PartialEq, DeriveActiveEnum)]
+#[sea_orm(rs_type = "i32", db_type = r#"Integer"#)]
+pub enum Color {
+    #[sea_orm(num_value = 0)]
+    Black,
+    #[sea_orm(num_value = 1)]
+    White,
+}
+
+#[derive(Debug, Clone, PartialEq, DeriveActiveEnum)]
+#[sea_orm(rs_type = "String", db_type = r#"Custom("tea".to_owned())"#)]
+pub enum Tea {
+    #[sea_orm(string_value = "EverydayTea")]
+    EverydayTea,
+    #[sea_orm(string_value = "BreakfastTea")]
+    BreakfastTea,
+}
diff --git a/tests/common/features/schema.rs b/tests/common/features/schema.rs
index 70752d35e..07a2953b4 100644
--- a/tests/common/features/schema.rs
+++ b/tests/common/features/schema.rs
@@ -2,8 +2,8 @@ pub use super::super::bakery_chain::*;
 
 use super::*;
 use crate::common::setup::create_table;
-use sea_orm::{error::*, sea_query, DatabaseConnection, DbConn, ExecResult};
-use sea_query::ColumnDef;
+use sea_orm::{ConnectionTrait, DatabaseConnection, DbConn, ExecResult, Statement, error::*, sea_query};
+use sea_query::{Alias, ColumnDef};
 
 pub async fn create_tables(db: &DatabaseConnection) -> Result<(), DbErr> {
     create_log_table(db).await?;
@@ -87,12 +87,9 @@ pub async fn create_active_enum_table(db: &DbConn) -> Result<ExecResult, DbErr>
                 .auto_increment()
                 .primary_key(),
         )
-        .col(
-            ColumnDef::new(active_enum::Column::Category)
-                .string_len(1)
-                .not_null(),
-        )
-        .col(ColumnDef::new(active_enum::Column::CategoryOpt).string_len(1))
+        .col(ColumnDef::new(active_enum::Column::Category).string_len(1))
+        .col(ColumnDef::new(active_enum::Column::Color).integer())
+        // .col(ColumnDef::new(active_enum::Column::Tea).custom(Alias::new("tea")))
         .to_owned();
 
     create_table(db, &stmt, ActiveEnum).await
diff --git a/tests/common/setup/mod.rs b/tests/common/setup/mod.rs
index bbe8baa77..615de234c 100644
--- a/tests/common/setup/mod.rs
+++ b/tests/common/setup/mod.rs
@@ -1,8 +1,8 @@
+use pretty_assertions::assert_eq;
 use sea_orm::{
     ConnectionTrait, Database, DatabaseBackend, DatabaseConnection, DbBackend, DbConn, DbErr,
     EntityTrait, ExecResult, Schema, Statement,
 };
-
 use sea_query::{Alias, Table, TableCreateStatement};
 
 pub async fn setup(base_url: &str, db_name: &str) -> DatabaseConnection {