Adds test_enum_extra_attributes

SeaQL · Nov 1, 2023 · 934b77a · 934b77a
1 parent e330a84
commit 934b77a
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 2 deletions.
diff --git a/sea-orm-codegen/src/entity/active_enum.rs b/sea-orm-codegen/src/entity/active_enum.rs
@@ -53,6 +53,7 @@ impl ActiveEnum {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::entity::writer::bonus_attributes;
     use pretty_assertions::assert_eq;
     use sea_query::{Alias, IntoIden};
 
@@ -111,4 +112,71 @@ mod tests {
             .to_string()
         )
     }
+
+    #[test]
+    fn test_enum_extra_attributes() {
+        assert_eq!(
+            ActiveEnum {
+                enum_name: Alias::new("coinflip_result_type").into_iden(),
+                values: vec!["HEADS", "TAILS"]
+                    .into_iter()
+                    .map(|variant| Alias::new(variant).into_iden())
+                    .collect(),
+            }
+            .impl_active_enum(
+                &WithSerde::None,
+                true,
+                &bonus_attributes([r#"serde(rename_all = "camelCase")"#])
+            )
+            .to_string(),
+            quote!(
+                #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
+                #[sea_orm(
+                    rs_type = "String",
+                    db_type = "Enum",
+                    enum_name = "coinflip_result_type"
+                )]
+                #[serde(rename_all = "camelCase")]
+                pub enum CoinflipResultType {
+                    #[sea_orm(string_value = "HEADS")]
+                    Heads,
+                    #[sea_orm(string_value = "TAILS")]
+                    Tails,
+                }
+            )
+            .to_string()
+        );
+        assert_eq!(
+            ActiveEnum {
+                enum_name: Alias::new("coinflip_result_type").into_iden(),
+                values: vec!["HEADS", "TAILS"]
+                    .into_iter()
+                    .map(|variant| Alias::new(variant).into_iden())
+                    .collect(),
+            }
+            .impl_active_enum(
+                &WithSerde::None,
+                true,
+                &bonus_attributes([r#"serde(rename_all = "camelCase")"#, "ts(export)"])
+            )
+            .to_string(),
+            quote!(
+                #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
+                #[sea_orm(
+                    rs_type = "String",
+                    db_type = "Enum",
+                    enum_name = "coinflip_result_type"
+                )]
+                #[serde(rename_all = "camelCase")]
+                #[ts(export)]
+                pub enum CoinflipResultType {
+                    #[sea_orm(string_value = "HEADS")]
+                    Heads,
+                    #[sea_orm(string_value = "TAILS")]
+                    Tails,
+                }
+            )
+            .to_string()
+        )
+    }
 }
diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs
@@ -95,8 +95,8 @@ where
         })
 }
 
-/// convert attributes argument to token stream
-fn bonus_attributes<T, I>(attributes: I) -> TokenStream
+/// convert *_extra_attributes argument to token stream
+pub(crate) fn bonus_attributes<T, I>(attributes: I) -> TokenStream
 where
     T: Into<String>,
     I: IntoIterator<Item = T>,