Fix duplicated active enum use statements on generated entities (#1953

) * fix: Only generate `use` statements once per enum * Adds gen_import_active_enum tests --------- Co-authored-by: Schlüter Lukas <[email protected]>
SeaQL · Nov 8, 2023 · 5fe81c5 · 5fe81c5
1 parent d8bf5ab
commit 5fe81c5
Showing 1 changed file with 144 additions and 10 deletions.
diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs
@@ -460,15 +460,23 @@ impl EntityWriter {
         entity
             .columns
             .iter()
-            .fold(TokenStream::new(), |mut ts, col| {
-                if let sea_query::ColumnType::Enum { name, .. } = col.get_inner_col_type() {
-                    let enum_name = format_ident!("{}", name.to_string().to_upper_camel_case());
-                    ts.extend([quote! {
-                        use super::sea_orm_active_enums::#enum_name;
-                    }]);
-                }
-                ts
-            })
+            .fold(
+                (TokenStream::new(), Vec::new()),
+                |(mut ts, mut enums), col| {
+                    if let sea_query::ColumnType::Enum { name, .. } = col.get_inner_col_type() {
+                        if !enums.contains(&name) {
+                            enums.push(name);
+                            let enum_name =
+                                format_ident!("{}", name.to_string().to_upper_camel_case());
+                            ts.extend([quote! {
+                                use super::sea_orm_active_enums::#enum_name;
+                            }]);
+                        }
+                    }
+                    (ts, enums)
+                },
+            )
+            .0
     }
 
     pub fn gen_model_struct(
@@ -814,7 +822,8 @@ mod tests {
     };
     use pretty_assertions::assert_eq;
     use proc_macro2::TokenStream;
-    use sea_query::{ColumnType, ForeignKeyAction, RcOrArc};
+    use quote::quote;
+    use sea_query::{Alias, ColumnType, ForeignKeyAction, RcOrArc, SeaRc};
     use std::io::{self, BufRead, BufReader, Read};
 
     fn setup() -> Vec<Entity> {
@@ -2275,4 +2284,129 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_gen_import_active_enum() -> io::Result<()> {
+        let entities = vec![
+            Entity {
+                table_name: "tea_pairing".to_owned(),
+                columns: vec![
+                    Column {
+                        name: "id".to_owned(),
+                        col_type: ColumnType::Integer,
+                        auto_increment: true,
+                        not_null: true,
+                        unique: false,
+                    },
+                    Column {
+                        name: "first_tea".to_owned(),
+                        col_type: ColumnType::Enum {
+                            name: SeaRc::new(Alias::new("tea_enum")),
+                            variants: vec![
+                                SeaRc::new(Alias::new("everyday_tea")),
+                                SeaRc::new(Alias::new("breakfast_tea")),
+                            ],
+                        },
+                        auto_increment: false,
+                        not_null: true,
+                        unique: false,
+                    },
+                    Column {
+                        name: "second_tea".to_owned(),
+                        col_type: ColumnType::Enum {
+                            name: SeaRc::new(Alias::new("tea_enum")),
+                            variants: vec![
+                                SeaRc::new(Alias::new("everyday_tea")),
+                                SeaRc::new(Alias::new("breakfast_tea")),
+                            ],
+                        },
+                        auto_increment: false,
+                        not_null: true,
+                        unique: false,
+                    },
+                ],
+                relations: vec![],
+                conjunct_relations: vec![],
+                primary_keys: vec![PrimaryKey {
+                    name: "id".to_owned(),
+                }],
+            },
+            Entity {
+                table_name: "tea_pairing_with_size".to_owned(),
+                columns: vec![
+                    Column {
+                        name: "id".to_owned(),
+                        col_type: ColumnType::Integer,
+                        auto_increment: true,
+                        not_null: true,
+                        unique: false,
+                    },
+                    Column {
+                        name: "first_tea".to_owned(),
+                        col_type: ColumnType::Enum {
+                            name: SeaRc::new(Alias::new("tea_enum")),
+                            variants: vec![
+                                SeaRc::new(Alias::new("everyday_tea")),
+                                SeaRc::new(Alias::new("breakfast_tea")),
+                            ],
+                        },
+                        auto_increment: false,
+                        not_null: true,
+                        unique: false,
+                    },
+                    Column {
+                        name: "second_tea".to_owned(),
+                        col_type: ColumnType::Enum {
+                            name: SeaRc::new(Alias::new("tea_enum")),
+                            variants: vec![
+                                SeaRc::new(Alias::new("everyday_tea")),
+                                SeaRc::new(Alias::new("breakfast_tea")),
+                            ],
+                        },
+                        auto_increment: false,
+                        not_null: true,
+                        unique: false,
+                    },
+                    Column {
+                        name: "size".to_owned(),
+                        col_type: ColumnType::Enum {
+                            name: SeaRc::new(Alias::new("tea_size")),
+                            variants: vec![
+                                SeaRc::new(Alias::new("small")),
+                                SeaRc::new(Alias::new("medium")),
+                                SeaRc::new(Alias::new("huge")),
+                            ],
+                        },
+                        auto_increment: false,
+                        not_null: true,
+                        unique: false,
+                    },
+                ],
+                relations: vec![],
+                conjunct_relations: vec![],
+                primary_keys: vec![PrimaryKey {
+                    name: "id".to_owned(),
+                }],
+            },
+        ];
+
+        assert_eq!(
+            quote!(
+                use super::sea_orm_active_enums::TeaEnum;
+            )
+            .to_string(),
+            EntityWriter::gen_import_active_enum(&entities[0]).to_string()
+        );
+
+        assert_eq!(
+            quote!(
+                use super::sea_orm_active_enums::TeaEnum;
+                use super::sea_orm_active_enums::TeaSize;
+            )
+            .to_string(),
+            EntityWriter::gen_import_active_enum(&entities[1]).to_string()
+        );
+
+        Ok(())
+    }
 }