sea_orm_codegen/entity/
writer.rs

1use crate::{util::escape_rust_keyword, ActiveEnum, Entity};
2use heck::ToUpperCamelCase;
3use proc_macro2::TokenStream;
4use quote::{format_ident, quote};
5use std::{collections::BTreeMap, str::FromStr};
6use syn::{punctuated::Punctuated, token::Comma};
7use tracing::info;
8
9#[derive(Clone, Debug)]
10pub struct EntityWriter {
11    pub(crate) entities: Vec<Entity>,
12    pub(crate) enums: BTreeMap<String, ActiveEnum>,
13}
14
15pub struct WriterOutput {
16    pub files: Vec<OutputFile>,
17}
18
19pub struct OutputFile {
20    pub name: String,
21    pub content: String,
22}
23
24#[derive(Default, PartialEq, Eq, Clone, Copy, Debug)]
25pub enum WithPrelude {
26    #[default]
27    All,
28    None,
29    AllAllowUnusedImports,
30}
31
32#[derive(PartialEq, Eq, Debug)]
33pub enum WithSerde {
34    None,
35    Serialize,
36    Deserialize,
37    Both,
38}
39
40#[derive(Debug)]
41pub enum DateTimeCrate {
42    Chrono,
43    Time,
44}
45
46#[derive(Debug)]
47pub struct EntityWriterContext {
48    pub(crate) expanded_format: bool,
49    pub(crate) frontend_format: bool,
50    pub(crate) with_prelude: WithPrelude,
51    pub(crate) with_serde: WithSerde,
52    pub(crate) with_copy_enums: bool,
53    pub(crate) date_time_crate: DateTimeCrate,
54    pub(crate) schema_name: Option<String>,
55    pub(crate) lib: bool,
56    pub(crate) serde_skip_hidden_column: bool,
57    pub(crate) serde_skip_deserializing_primary_key: bool,
58    pub(crate) model_extra_derives: TokenStream,
59    pub(crate) model_extra_attributes: TokenStream,
60    pub(crate) enum_extra_derives: TokenStream,
61    pub(crate) enum_extra_attributes: TokenStream,
62    pub(crate) seaography: bool,
63    pub(crate) impl_active_model_behavior: bool,
64}
65
66impl WithSerde {
67    pub fn extra_derive(&self) -> TokenStream {
68        let mut extra_derive = match self {
69            Self::None => {
70                quote! {}
71            }
72            Self::Serialize => {
73                quote! {
74                    Serialize
75                }
76            }
77            Self::Deserialize => {
78                quote! {
79                    Deserialize
80                }
81            }
82            Self::Both => {
83                quote! {
84                    Serialize, Deserialize
85                }
86            }
87        };
88        if !extra_derive.is_empty() {
89            extra_derive = quote! { , #extra_derive }
90        }
91        extra_derive
92    }
93}
94
95/// Converts *_extra_derives argument to token stream
96pub(crate) fn bonus_derive<T, I>(extra_derives: I) -> TokenStream
97where
98    T: Into<String>,
99    I: IntoIterator<Item = T>,
100{
101    extra_derives.into_iter().map(Into::<String>::into).fold(
102        TokenStream::default(),
103        |acc, derive| {
104            let tokens: TokenStream = derive.parse().unwrap();
105            quote! { #acc, #tokens }
106        },
107    )
108}
109
110/// convert *_extra_attributes argument to token stream
111pub(crate) fn bonus_attributes<T, I>(attributes: I) -> TokenStream
112where
113    T: Into<String>,
114    I: IntoIterator<Item = T>,
115{
116    attributes.into_iter().map(Into::<String>::into).fold(
117        TokenStream::default(),
118        |acc, attribute| {
119            let tokens: TokenStream = attribute.parse().unwrap();
120            quote! {
121                #acc
122                #[#tokens]
123            }
124        },
125    )
126}
127
128impl FromStr for WithPrelude {
129    type Err = crate::Error;
130
131    fn from_str(s: &str) -> Result<Self, Self::Err> {
132        Ok(match s {
133            "none" => Self::None,
134            "all-allow-unused-imports" => Self::AllAllowUnusedImports,
135            "all" => Self::All,
136            v => {
137                return Err(crate::Error::TransformError(format!(
138                    "Unsupported enum variant '{v}'"
139                )))
140            }
141        })
142    }
143}
144
145impl FromStr for WithSerde {
146    type Err = crate::Error;
147
148    fn from_str(s: &str) -> Result<Self, Self::Err> {
149        Ok(match s {
150            "none" => Self::None,
151            "serialize" => Self::Serialize,
152            "deserialize" => Self::Deserialize,
153            "both" => Self::Both,
154            v => {
155                return Err(crate::Error::TransformError(format!(
156                    "Unsupported enum variant '{v}'"
157                )))
158            }
159        })
160    }
161}
162
163impl EntityWriterContext {
164    #[allow(clippy::too_many_arguments)]
165    pub fn new(
166        expanded_format: bool,
167        frontend_format: bool,
168        with_prelude: WithPrelude,
169        with_serde: WithSerde,
170        with_copy_enums: bool,
171        date_time_crate: DateTimeCrate,
172        schema_name: Option<String>,
173        lib: bool,
174        serde_skip_deserializing_primary_key: bool,
175        serde_skip_hidden_column: bool,
176        model_extra_derives: Vec<String>,
177        model_extra_attributes: Vec<String>,
178        enum_extra_derives: Vec<String>,
179        enum_extra_attributes: Vec<String>,
180        seaography: bool,
181        impl_active_model_behavior: bool,
182    ) -> Self {
183        Self {
184            expanded_format,
185            frontend_format,
186            with_prelude,
187            with_serde,
188            with_copy_enums,
189            date_time_crate,
190            schema_name,
191            lib,
192            serde_skip_deserializing_primary_key,
193            serde_skip_hidden_column,
194            model_extra_derives: bonus_derive(model_extra_derives),
195            model_extra_attributes: bonus_attributes(model_extra_attributes),
196            enum_extra_derives: bonus_derive(enum_extra_derives),
197            enum_extra_attributes: bonus_attributes(enum_extra_attributes),
198            seaography,
199            impl_active_model_behavior,
200        }
201    }
202}
203
204impl EntityWriter {
205    pub fn generate(self, context: &EntityWriterContext) -> WriterOutput {
206        let mut files = Vec::new();
207        files.extend(self.write_entities(context));
208        let with_prelude = context.with_prelude != WithPrelude::None;
209        files.push(self.write_index_file(context.lib, with_prelude, context.seaography));
210        if with_prelude {
211            files.push(self.write_prelude(context.with_prelude, context.frontend_format));
212        }
213        if !self.enums.is_empty() {
214            files.push(self.write_sea_orm_active_enums(
215                &context.with_serde,
216                context.with_copy_enums,
217                &context.enum_extra_derives,
218                &context.enum_extra_attributes,
219                context.frontend_format,
220            ));
221        }
222        WriterOutput { files }
223    }
224
225    pub fn write_entities(&self, context: &EntityWriterContext) -> Vec<OutputFile> {
226        self.entities
227            .iter()
228            .map(|entity| {
229                let entity_file = format!("{}.rs", entity.get_table_name_snake_case());
230                let column_info = entity
231                    .columns
232                    .iter()
233                    .map(|column| column.get_info(&context.date_time_crate))
234                    .collect::<Vec<String>>();
235                // Serde must be enabled to use this
236                let serde_skip_deserializing_primary_key = context
237                    .serde_skip_deserializing_primary_key
238                    && matches!(context.with_serde, WithSerde::Both | WithSerde::Deserialize);
239                let serde_skip_hidden_column = context.serde_skip_hidden_column
240                    && matches!(
241                        context.with_serde,
242                        WithSerde::Both | WithSerde::Serialize | WithSerde::Deserialize
243                    );
244
245                info!("Generating {}", entity_file);
246                for info in column_info.iter() {
247                    info!("    > {}", info);
248                }
249
250                let mut lines = Vec::new();
251                Self::write_doc_comment(&mut lines);
252                let code_blocks = if context.frontend_format {
253                    Self::gen_frontend_code_blocks(
254                        entity,
255                        &context.with_serde,
256                        &context.date_time_crate,
257                        &context.schema_name,
258                        serde_skip_deserializing_primary_key,
259                        serde_skip_hidden_column,
260                        &context.model_extra_derives,
261                        &context.model_extra_attributes,
262                        context.seaography,
263                        context.impl_active_model_behavior,
264                    )
265                } else if context.expanded_format {
266                    Self::gen_expanded_code_blocks(
267                        entity,
268                        &context.with_serde,
269                        &context.date_time_crate,
270                        &context.schema_name,
271                        serde_skip_deserializing_primary_key,
272                        serde_skip_hidden_column,
273                        &context.model_extra_derives,
274                        &context.model_extra_attributes,
275                        context.seaography,
276                        context.impl_active_model_behavior,
277                    )
278                } else {
279                    Self::gen_compact_code_blocks(
280                        entity,
281                        &context.with_serde,
282                        &context.date_time_crate,
283                        &context.schema_name,
284                        serde_skip_deserializing_primary_key,
285                        serde_skip_hidden_column,
286                        &context.model_extra_derives,
287                        &context.model_extra_attributes,
288                        context.seaography,
289                        context.impl_active_model_behavior,
290                    )
291                };
292                Self::write(&mut lines, code_blocks);
293                OutputFile {
294                    name: entity_file,
295                    content: lines.join("\n\n"),
296                }
297            })
298            .collect()
299    }
300
301    pub fn write_index_file(&self, lib: bool, prelude: bool, seaography: bool) -> OutputFile {
302        let mut lines = Vec::new();
303        Self::write_doc_comment(&mut lines);
304        let code_blocks: Vec<TokenStream> = self.entities.iter().map(Self::gen_mod).collect();
305        if prelude {
306            Self::write(
307                &mut lines,
308                vec![quote! {
309                    pub mod prelude;
310                }],
311            );
312            lines.push("".to_owned());
313        }
314        Self::write(&mut lines, code_blocks);
315        if !self.enums.is_empty() {
316            Self::write(
317                &mut lines,
318                vec![quote! {
319                    pub mod sea_orm_active_enums;
320                }],
321            );
322        }
323
324        if seaography {
325            lines.push("".to_owned());
326            let ts = Self::gen_seaography_entity_mod(&self.entities, &self.enums);
327            Self::write(&mut lines, vec![ts]);
328        }
329
330        let file_name = match lib {
331            true => "lib.rs".to_owned(),
332            false => "mod.rs".to_owned(),
333        };
334
335        OutputFile {
336            name: file_name,
337            content: lines.join("\n"),
338        }
339    }
340
341    pub fn write_prelude(&self, with_prelude: WithPrelude, frontend_format: bool) -> OutputFile {
342        let mut lines = Vec::new();
343        Self::write_doc_comment(&mut lines);
344        if with_prelude == WithPrelude::AllAllowUnusedImports {
345            Self::write_allow_unused_imports(&mut lines)
346        }
347        let code_blocks = self
348            .entities
349            .iter()
350            .map({
351                if frontend_format {
352                    Self::gen_prelude_use_model
353                } else {
354                    Self::gen_prelude_use
355                }
356            })
357            .collect();
358        Self::write(&mut lines, code_blocks);
359        OutputFile {
360            name: "prelude.rs".to_owned(),
361            content: lines.join("\n"),
362        }
363    }
364
365    pub fn write_sea_orm_active_enums(
366        &self,
367        with_serde: &WithSerde,
368        with_copy_enums: bool,
369        extra_derives: &TokenStream,
370        extra_attributes: &TokenStream,
371        frontend_format: bool,
372    ) -> OutputFile {
373        let mut lines = Vec::new();
374        Self::write_doc_comment(&mut lines);
375        if frontend_format {
376            Self::write(&mut lines, vec![Self::gen_import_serde(with_serde)]);
377        } else {
378            Self::write(&mut lines, vec![Self::gen_import(with_serde)]);
379        }
380        lines.push("".to_owned());
381        let code_blocks = self
382            .enums
383            .values()
384            .map(|active_enum| {
385                active_enum.impl_active_enum(
386                    with_serde,
387                    with_copy_enums,
388                    extra_derives,
389                    extra_attributes,
390                    frontend_format,
391                )
392            })
393            .collect();
394        Self::write(&mut lines, code_blocks);
395        OutputFile {
396            name: "sea_orm_active_enums.rs".to_owned(),
397            content: lines.join("\n"),
398        }
399    }
400
401    pub fn write(lines: &mut Vec<String>, code_blocks: Vec<TokenStream>) {
402        lines.extend(
403            code_blocks
404                .into_iter()
405                .map(|code_block| code_block.to_string())
406                .collect::<Vec<_>>(),
407        );
408    }
409
410    pub fn write_doc_comment(lines: &mut Vec<String>) {
411        let ver = env!("CARGO_PKG_VERSION");
412        let comments = vec![format!(
413            "//! `SeaORM` Entity, @generated by sea-orm-codegen {ver}"
414        )];
415        lines.extend(comments);
416        lines.push("".to_owned());
417    }
418
419    pub fn write_allow_unused_imports(lines: &mut Vec<String>) {
420        lines.extend(vec!["#![allow(unused_imports)]".to_string()]);
421        lines.push("".to_owned());
422    }
423
424    #[allow(clippy::too_many_arguments)]
425    pub fn gen_expanded_code_blocks(
426        entity: &Entity,
427        with_serde: &WithSerde,
428        date_time_crate: &DateTimeCrate,
429        schema_name: &Option<String>,
430        serde_skip_deserializing_primary_key: bool,
431        serde_skip_hidden_column: bool,
432        model_extra_derives: &TokenStream,
433        model_extra_attributes: &TokenStream,
434        seaography: bool,
435        impl_active_model_behavior: bool,
436    ) -> Vec<TokenStream> {
437        let mut imports = Self::gen_import(with_serde);
438        imports.extend(Self::gen_import_active_enum(entity));
439        let mut code_blocks = vec![
440            imports,
441            Self::gen_entity_struct(),
442            Self::gen_impl_entity_name(entity, schema_name),
443            Self::gen_model_struct(
444                entity,
445                with_serde,
446                date_time_crate,
447                serde_skip_deserializing_primary_key,
448                serde_skip_hidden_column,
449                model_extra_derives,
450                model_extra_attributes,
451            ),
452            Self::gen_column_enum(entity),
453            Self::gen_primary_key_enum(entity),
454            Self::gen_impl_primary_key(entity, date_time_crate),
455            Self::gen_relation_enum(entity),
456            Self::gen_impl_column_trait(entity),
457            Self::gen_impl_relation_trait(entity),
458        ];
459        code_blocks.extend(Self::gen_impl_related(entity));
460        code_blocks.extend(Self::gen_impl_conjunct_related(entity));
461        if impl_active_model_behavior {
462            code_blocks.extend([Self::impl_active_model_behavior()]);
463        }
464        if seaography {
465            code_blocks.extend([Self::gen_related_entity(entity)]);
466        }
467        code_blocks
468    }
469
470    #[allow(clippy::too_many_arguments)]
471    pub fn gen_compact_code_blocks(
472        entity: &Entity,
473        with_serde: &WithSerde,
474        date_time_crate: &DateTimeCrate,
475        schema_name: &Option<String>,
476        serde_skip_deserializing_primary_key: bool,
477        serde_skip_hidden_column: bool,
478        model_extra_derives: &TokenStream,
479        model_extra_attributes: &TokenStream,
480        seaography: bool,
481        impl_active_model_behavior: bool,
482    ) -> Vec<TokenStream> {
483        let mut imports = Self::gen_import(with_serde);
484        imports.extend(Self::gen_import_active_enum(entity));
485        let mut code_blocks = vec![
486            imports,
487            Self::gen_compact_model_struct(
488                entity,
489                with_serde,
490                date_time_crate,
491                schema_name,
492                serde_skip_deserializing_primary_key,
493                serde_skip_hidden_column,
494                model_extra_derives,
495                model_extra_attributes,
496            ),
497            Self::gen_compact_relation_enum(entity),
498        ];
499        code_blocks.extend(Self::gen_impl_related(entity));
500        code_blocks.extend(Self::gen_impl_conjunct_related(entity));
501        if impl_active_model_behavior {
502            code_blocks.extend([Self::impl_active_model_behavior()]);
503        }
504        if seaography {
505            code_blocks.extend([Self::gen_related_entity(entity)]);
506        }
507        code_blocks
508    }
509
510    #[allow(clippy::too_many_arguments)]
511    pub fn gen_frontend_code_blocks(
512        entity: &Entity,
513        with_serde: &WithSerde,
514        date_time_crate: &DateTimeCrate,
515        schema_name: &Option<String>,
516        serde_skip_deserializing_primary_key: bool,
517        serde_skip_hidden_column: bool,
518        model_extra_derives: &TokenStream,
519        model_extra_attributes: &TokenStream,
520        _seaography: bool,
521        _impl_active_model_behavior: bool,
522    ) -> Vec<TokenStream> {
523        let mut imports = Self::gen_import_serde(with_serde);
524        imports.extend(Self::gen_import_active_enum(entity));
525        let code_blocks = vec![
526            imports,
527            Self::gen_frontend_model_struct(
528                entity,
529                with_serde,
530                date_time_crate,
531                schema_name,
532                serde_skip_deserializing_primary_key,
533                serde_skip_hidden_column,
534                model_extra_derives,
535                model_extra_attributes,
536            ),
537        ];
538        code_blocks
539    }
540
541    pub fn gen_import(with_serde: &WithSerde) -> TokenStream {
542        let serde_import = Self::gen_import_serde(with_serde);
543        quote! {
544            use sea_orm::entity::prelude::*;
545            #serde_import
546        }
547    }
548
549    pub fn gen_import_serde(with_serde: &WithSerde) -> TokenStream {
550        match with_serde {
551            WithSerde::None => Default::default(),
552            WithSerde::Serialize => {
553                quote! {
554                    use serde::Serialize;
555                }
556            }
557            WithSerde::Deserialize => {
558                quote! {
559                    use serde::Deserialize;
560                }
561            }
562            WithSerde::Both => {
563                quote! {
564                    use serde::{Deserialize,Serialize};
565                }
566            }
567        }
568    }
569
570    pub fn gen_entity_struct() -> TokenStream {
571        quote! {
572            #[derive(Copy, Clone, Default, Debug, DeriveEntity)]
573            pub struct Entity;
574        }
575    }
576
577    pub fn gen_impl_entity_name(entity: &Entity, schema_name: &Option<String>) -> TokenStream {
578        let schema_name = match Self::gen_schema_name(schema_name) {
579            Some(schema_name) => quote! {
580                fn schema_name(&self) -> Option<&str> {
581                    Some(#schema_name)
582                }
583            },
584            None => quote! {},
585        };
586        let table_name = entity.table_name.as_str();
587        let table_name = quote! {
588            fn table_name(&self) -> &str {
589                #table_name
590            }
591        };
592        quote! {
593            impl EntityName for Entity {
594                #schema_name
595                #table_name
596            }
597        }
598    }
599
600    pub fn gen_import_active_enum(entity: &Entity) -> TokenStream {
601        entity
602            .columns
603            .iter()
604            .fold(
605                (TokenStream::new(), Vec::new()),
606                |(mut ts, mut enums), col| {
607                    if let sea_query::ColumnType::Enum { name, .. } = col.get_inner_col_type() {
608                        if !enums.contains(&name) {
609                            enums.push(name);
610                            let enum_name =
611                                format_ident!("{}", name.to_string().to_upper_camel_case());
612                            ts.extend([quote! {
613                                use super::sea_orm_active_enums::#enum_name;
614                            }]);
615                        }
616                    }
617                    (ts, enums)
618                },
619            )
620            .0
621    }
622
623    pub fn gen_model_struct(
624        entity: &Entity,
625        with_serde: &WithSerde,
626        date_time_crate: &DateTimeCrate,
627        serde_skip_deserializing_primary_key: bool,
628        serde_skip_hidden_column: bool,
629        model_extra_derives: &TokenStream,
630        model_extra_attributes: &TokenStream,
631    ) -> TokenStream {
632        let column_names_snake_case = entity.get_column_names_snake_case();
633        let column_rs_types = entity.get_column_rs_types(date_time_crate);
634        let if_eq_needed = entity.get_eq_needed();
635        let serde_attributes = entity.get_column_serde_attributes(
636            serde_skip_deserializing_primary_key,
637            serde_skip_hidden_column,
638        );
639        let extra_derive = with_serde.extra_derive();
640
641        quote! {
642            #[derive(Clone, Debug, PartialEq, DeriveModel, DeriveActiveModel #if_eq_needed #extra_derive #model_extra_derives)]
643            #model_extra_attributes
644            pub struct Model {
645                #(
646                    #serde_attributes
647                    pub #column_names_snake_case: #column_rs_types,
648                )*
649            }
650        }
651    }
652
653    pub fn gen_column_enum(entity: &Entity) -> TokenStream {
654        let column_variants = entity.columns.iter().map(|col| {
655            let variant = col.get_name_camel_case();
656            let mut variant = quote! { #variant };
657            if !col.is_snake_case_name() {
658                let column_name = &col.name;
659                variant = quote! {
660                    #[sea_orm(column_name = #column_name)]
661                    #variant
662                };
663            }
664            variant
665        });
666        quote! {
667            #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
668            pub enum Column {
669                #(#column_variants,)*
670            }
671        }
672    }
673
674    pub fn gen_primary_key_enum(entity: &Entity) -> TokenStream {
675        let primary_key_names_camel_case = entity.get_primary_key_names_camel_case();
676        quote! {
677            #[derive(Copy, Clone, Debug, EnumIter, DerivePrimaryKey)]
678            pub enum PrimaryKey {
679                #(#primary_key_names_camel_case,)*
680            }
681        }
682    }
683
684    pub fn gen_impl_primary_key(entity: &Entity, date_time_crate: &DateTimeCrate) -> TokenStream {
685        let primary_key_auto_increment = entity.get_primary_key_auto_increment();
686        let value_type = entity.get_primary_key_rs_type(date_time_crate);
687        quote! {
688            impl PrimaryKeyTrait for PrimaryKey {
689                type ValueType = #value_type;
690
691                fn auto_increment() -> bool {
692                    #primary_key_auto_increment
693                }
694            }
695        }
696    }
697
698    pub fn gen_relation_enum(entity: &Entity) -> TokenStream {
699        let relation_enum_name = entity.get_relation_enum_name();
700        quote! {
701            #[derive(Copy, Clone, Debug, EnumIter)]
702            pub enum Relation {
703                #(#relation_enum_name,)*
704            }
705        }
706    }
707
708    pub fn gen_impl_column_trait(entity: &Entity) -> TokenStream {
709        let column_names_camel_case = entity.get_column_names_camel_case();
710        let column_defs = entity.get_column_defs();
711        quote! {
712            impl ColumnTrait for Column {
713                type EntityName = Entity;
714
715                fn def(&self) -> ColumnDef {
716                    match self {
717                        #(Self::#column_names_camel_case => #column_defs,)*
718                    }
719                }
720            }
721        }
722    }
723
724    pub fn gen_impl_relation_trait(entity: &Entity) -> TokenStream {
725        let relation_enum_name = entity.get_relation_enum_name();
726        let relation_defs = entity.get_relation_defs();
727        let quoted = if relation_enum_name.is_empty() {
728            quote! {
729                panic!("No RelationDef")
730            }
731        } else {
732            quote! {
733                match self {
734                    #(Self::#relation_enum_name => #relation_defs,)*
735                }
736            }
737        };
738        quote! {
739            impl RelationTrait for Relation {
740                fn def(&self) -> RelationDef {
741                    #quoted
742                }
743            }
744        }
745    }
746
747    pub fn gen_impl_related(entity: &Entity) -> Vec<TokenStream> {
748        entity
749            .relations
750            .iter()
751            .filter(|rel| !rel.self_referencing && rel.num_suffix == 0 && rel.impl_related)
752            .map(|rel| {
753                let enum_name = rel.get_enum_name();
754                let module_name = rel.get_module_name();
755                let inner = quote! {
756                    fn to() -> RelationDef {
757                        Relation::#enum_name.def()
758                    }
759                };
760                if module_name.is_some() {
761                    quote! {
762                        impl Related<super::#module_name::Entity> for Entity { #inner }
763                    }
764                } else {
765                    quote! {
766                        impl Related<Entity> for Entity { #inner }
767                    }
768                }
769            })
770            .collect()
771    }
772
773    /// Used to generate `enum RelatedEntity` that is useful to the Seaography project
774    pub fn gen_related_entity(entity: &Entity) -> TokenStream {
775        let related_enum_name = entity.get_related_entity_enum_name();
776        let related_attrs = entity.get_related_entity_attrs();
777
778        quote! {
779            #[derive(Copy, Clone, Debug, EnumIter, DeriveRelatedEntity)]
780            pub enum RelatedEntity {
781                #(
782                    #related_attrs
783                    #related_enum_name
784                ),*
785            }
786        }
787    }
788
789    pub fn gen_impl_conjunct_related(entity: &Entity) -> Vec<TokenStream> {
790        let table_name_camel_case = entity.get_table_name_camel_case_ident();
791        let via_snake_case = entity.get_conjunct_relations_via_snake_case();
792        let to_snake_case = entity.get_conjunct_relations_to_snake_case();
793        let to_upper_camel_case = entity.get_conjunct_relations_to_upper_camel_case();
794        via_snake_case
795            .into_iter()
796            .zip(to_snake_case)
797            .zip(to_upper_camel_case)
798            .map(|((via_snake_case, to_snake_case), to_upper_camel_case)| {
799                quote! {
800                    impl Related<super::#to_snake_case::Entity> for Entity {
801                        fn to() -> RelationDef {
802                            super::#via_snake_case::Relation::#to_upper_camel_case.def()
803                        }
804
805                        fn via() -> Option<RelationDef> {
806                            Some(super::#via_snake_case::Relation::#table_name_camel_case.def().rev())
807                        }
808                    }
809                }
810            })
811            .collect()
812    }
813
814    pub fn impl_active_model_behavior() -> TokenStream {
815        quote! {
816            impl ActiveModelBehavior for ActiveModel {}
817        }
818    }
819
820    pub fn gen_mod(entity: &Entity) -> TokenStream {
821        let table_name_snake_case_ident = format_ident!(
822            "{}",
823            escape_rust_keyword(entity.get_table_name_snake_case_ident())
824        );
825        quote! {
826            pub mod #table_name_snake_case_ident;
827        }
828    }
829
830    pub fn gen_seaography_entity_mod(
831        entities: &[Entity],
832        enums: &BTreeMap<String, ActiveEnum>,
833    ) -> TokenStream {
834        let mut ts = TokenStream::new();
835        for entity in entities {
836            let table_name_snake_case_ident = format_ident!(
837                "{}",
838                escape_rust_keyword(entity.get_table_name_snake_case_ident())
839            );
840            ts = quote! {
841                #ts
842                #table_name_snake_case_ident,
843            }
844        }
845        ts = quote! {
846            seaography::register_entity_modules!([
847                #ts
848            ]);
849        };
850
851        let mut enum_ts = TokenStream::new();
852        for active_enum in enums.values() {
853            let enum_name = &active_enum.enum_name.to_string();
854            let enum_iden = format_ident!("{}", enum_name.to_upper_camel_case());
855            enum_ts = quote! {
856                #enum_ts
857                sea_orm_active_enums::#enum_iden,
858            }
859        }
860        if !enum_ts.is_empty() {
861            ts = quote! {
862                #ts
863
864                seaography::register_active_enums!([
865                    #enum_ts
866                ]);
867            };
868        }
869        ts
870    }
871
872    pub fn gen_prelude_use(entity: &Entity) -> TokenStream {
873        let table_name_snake_case_ident = entity.get_table_name_snake_case_ident();
874        let table_name_camel_case_ident = entity.get_table_name_camel_case_ident();
875        quote! {
876            pub use super::#table_name_snake_case_ident::Entity as #table_name_camel_case_ident;
877        }
878    }
879
880    pub fn gen_prelude_use_model(entity: &Entity) -> TokenStream {
881        let table_name_snake_case_ident = entity.get_table_name_snake_case_ident();
882        let table_name_camel_case_ident = entity.get_table_name_camel_case_ident();
883        quote! {
884            pub use super::#table_name_snake_case_ident::Model as #table_name_camel_case_ident;
885        }
886    }
887
888    #[allow(clippy::too_many_arguments)]
889    pub fn gen_compact_model_struct(
890        entity: &Entity,
891        with_serde: &WithSerde,
892        date_time_crate: &DateTimeCrate,
893        schema_name: &Option<String>,
894        serde_skip_deserializing_primary_key: bool,
895        serde_skip_hidden_column: bool,
896        model_extra_derives: &TokenStream,
897        model_extra_attributes: &TokenStream,
898    ) -> TokenStream {
899        let table_name = entity.table_name.as_str();
900        let column_names_snake_case = entity.get_column_names_snake_case();
901        let column_rs_types = entity.get_column_rs_types(date_time_crate);
902        let if_eq_needed = entity.get_eq_needed();
903        let primary_keys: Vec<String> = entity
904            .primary_keys
905            .iter()
906            .map(|pk| pk.name.clone())
907            .collect();
908        let attrs: Vec<TokenStream> = entity
909            .columns
910            .iter()
911            .map(|col| {
912                let mut attrs: Punctuated<_, Comma> = Punctuated::new();
913                let is_primary_key = primary_keys.contains(&col.name);
914                if !col.is_snake_case_name() {
915                    let column_name = &col.name;
916                    attrs.push(quote! { column_name = #column_name });
917                }
918                if is_primary_key {
919                    attrs.push(quote! { primary_key });
920                    if !col.auto_increment {
921                        attrs.push(quote! { auto_increment = false });
922                    }
923                }
924                if let Some(ts) = col.get_col_type_attrs() {
925                    attrs.extend([ts]);
926                    if !col.not_null {
927                        attrs.push(quote! { nullable });
928                    }
929                };
930                if col.unique {
931                    attrs.push(quote! { unique });
932                }
933                let mut ts = quote! {};
934                if !attrs.is_empty() {
935                    for (i, attr) in attrs.into_iter().enumerate() {
936                        if i > 0 {
937                            ts = quote! { #ts, };
938                        }
939                        ts = quote! { #ts #attr };
940                    }
941                    ts = quote! { #[sea_orm(#ts)] };
942                }
943                let serde_attribute = col.get_serde_attribute(
944                    is_primary_key,
945                    serde_skip_deserializing_primary_key,
946                    serde_skip_hidden_column,
947                );
948                ts = quote! {
949                    #ts
950                    #serde_attribute
951                };
952                ts
953            })
954            .collect();
955        let schema_name = match Self::gen_schema_name(schema_name) {
956            Some(schema_name) => quote! {
957                schema_name = #schema_name,
958            },
959            None => quote! {},
960        };
961        let extra_derive = with_serde.extra_derive();
962
963        quote! {
964            #[derive(Clone, Debug, PartialEq, DeriveEntityModel #if_eq_needed #extra_derive #model_extra_derives)]
965            #[sea_orm(
966                #schema_name
967                table_name = #table_name
968            )]
969            #model_extra_attributes
970            pub struct Model {
971                #(
972                    #attrs
973                    pub #column_names_snake_case: #column_rs_types,
974                )*
975            }
976        }
977    }
978
979    pub fn gen_compact_relation_enum(entity: &Entity) -> TokenStream {
980        let relation_enum_name = entity.get_relation_enum_name();
981        let attrs = entity.get_relation_attrs();
982        quote! {
983            #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
984            pub enum Relation {
985                #(
986                    #attrs
987                    #relation_enum_name,
988                )*
989            }
990        }
991    }
992
993    #[allow(clippy::too_many_arguments)]
994    pub fn gen_frontend_model_struct(
995        entity: &Entity,
996        with_serde: &WithSerde,
997        date_time_crate: &DateTimeCrate,
998        _schema_name: &Option<String>,
999        serde_skip_deserializing_primary_key: bool,
1000        serde_skip_hidden_column: bool,
1001        model_extra_derives: &TokenStream,
1002        model_extra_attributes: &TokenStream,
1003    ) -> TokenStream {
1004        let column_names_snake_case = entity.get_column_names_snake_case();
1005        let column_rs_types = entity.get_column_rs_types(date_time_crate);
1006        let if_eq_needed = entity.get_eq_needed();
1007        let primary_keys: Vec<String> = entity
1008            .primary_keys
1009            .iter()
1010            .map(|pk| pk.name.clone())
1011            .collect();
1012        let attrs: Vec<TokenStream> = entity
1013            .columns
1014            .iter()
1015            .map(|col| {
1016                let is_primary_key = primary_keys.contains(&col.name);
1017                col.get_serde_attribute(
1018                    is_primary_key,
1019                    serde_skip_deserializing_primary_key,
1020                    serde_skip_hidden_column,
1021                )
1022            })
1023            .collect();
1024        let extra_derive = with_serde.extra_derive();
1025
1026        quote! {
1027            #[derive(Clone, Debug, PartialEq #if_eq_needed #extra_derive #model_extra_derives)]
1028            #model_extra_attributes
1029            pub struct Model {
1030                #(
1031                    #attrs
1032                    pub #column_names_snake_case: #column_rs_types,
1033                )*
1034            }
1035        }
1036    }
1037
1038    pub fn gen_schema_name(schema_name: &Option<String>) -> Option<TokenStream> {
1039        schema_name
1040            .as_ref()
1041            .map(|schema_name| quote! { #schema_name })
1042    }
1043}
1044
1045#[cfg(test)]
1046mod tests {
1047    use crate::{
1048        entity::writer::{bonus_attributes, bonus_derive},
1049        Column, ConjunctRelation, DateTimeCrate, Entity, EntityWriter, PrimaryKey, Relation,
1050        RelationType, WithSerde,
1051    };
1052    use pretty_assertions::assert_eq;
1053    use proc_macro2::TokenStream;
1054    use quote::quote;
1055    use sea_query::{Alias, ColumnType, ForeignKeyAction, RcOrArc, SeaRc, StringLen};
1056    use std::io::{self, BufRead, BufReader, Read};
1057
1058    fn setup() -> Vec<Entity> {
1059        vec![
1060            Entity {
1061                table_name: "cake".to_owned(),
1062                columns: vec![
1063                    Column {
1064                        name: "id".to_owned(),
1065                        col_type: ColumnType::Integer,
1066                        auto_increment: true,
1067                        not_null: true,
1068                        unique: false,
1069                    },
1070                    Column {
1071                        name: "name".to_owned(),
1072                        col_type: ColumnType::Text,
1073                        auto_increment: false,
1074                        not_null: false,
1075                        unique: false,
1076                    },
1077                ],
1078                relations: vec![Relation {
1079                    ref_table: "fruit".to_owned(),
1080                    columns: vec![],
1081                    ref_columns: vec![],
1082                    rel_type: RelationType::HasMany,
1083                    on_delete: None,
1084                    on_update: None,
1085                    self_referencing: false,
1086                    num_suffix: 0,
1087                    impl_related: true,
1088                }],
1089                conjunct_relations: vec![ConjunctRelation {
1090                    via: "cake_filling".to_owned(),
1091                    to: "filling".to_owned(),
1092                }],
1093                primary_keys: vec![PrimaryKey {
1094                    name: "id".to_owned(),
1095                }],
1096            },
1097            Entity {
1098                table_name: "_cake_filling_".to_owned(),
1099                columns: vec![
1100                    Column {
1101                        name: "cake_id".to_owned(),
1102                        col_type: ColumnType::Integer,
1103                        auto_increment: false,
1104                        not_null: true,
1105                        unique: false,
1106                    },
1107                    Column {
1108                        name: "filling_id".to_owned(),
1109                        col_type: ColumnType::Integer,
1110                        auto_increment: false,
1111                        not_null: true,
1112                        unique: false,
1113                    },
1114                ],
1115                relations: vec![
1116                    Relation {
1117                        ref_table: "cake".to_owned(),
1118                        columns: vec!["cake_id".to_owned()],
1119                        ref_columns: vec!["id".to_owned()],
1120                        rel_type: RelationType::BelongsTo,
1121                        on_delete: Some(ForeignKeyAction::Cascade),
1122                        on_update: Some(ForeignKeyAction::Cascade),
1123                        self_referencing: false,
1124                        num_suffix: 0,
1125                        impl_related: true,
1126                    },
1127                    Relation {
1128                        ref_table: "filling".to_owned(),
1129                        columns: vec!["filling_id".to_owned()],
1130                        ref_columns: vec!["id".to_owned()],
1131                        rel_type: RelationType::BelongsTo,
1132                        on_delete: Some(ForeignKeyAction::Cascade),
1133                        on_update: Some(ForeignKeyAction::Cascade),
1134                        self_referencing: false,
1135                        num_suffix: 0,
1136                        impl_related: true,
1137                    },
1138                ],
1139                conjunct_relations: vec![],
1140                primary_keys: vec![
1141                    PrimaryKey {
1142                        name: "cake_id".to_owned(),
1143                    },
1144                    PrimaryKey {
1145                        name: "filling_id".to_owned(),
1146                    },
1147                ],
1148            },
1149            Entity {
1150                table_name: "cake_filling_price".to_owned(),
1151                columns: vec![
1152                    Column {
1153                        name: "cake_id".to_owned(),
1154                        col_type: ColumnType::Integer,
1155                        auto_increment: false,
1156                        not_null: true,
1157                        unique: false,
1158                    },
1159                    Column {
1160                        name: "filling_id".to_owned(),
1161                        col_type: ColumnType::Integer,
1162                        auto_increment: false,
1163                        not_null: true,
1164                        unique: false,
1165                    },
1166                    Column {
1167                        name: "price".to_owned(),
1168                        col_type: ColumnType::Decimal(None),
1169                        auto_increment: false,
1170                        not_null: true,
1171                        unique: false,
1172                    },
1173                ],
1174                relations: vec![Relation {
1175                    ref_table: "cake_filling".to_owned(),
1176                    columns: vec!["cake_id".to_owned(), "filling_id".to_owned()],
1177                    ref_columns: vec!["cake_id".to_owned(), "filling_id".to_owned()],
1178                    rel_type: RelationType::BelongsTo,
1179                    on_delete: None,
1180                    on_update: None,
1181                    self_referencing: false,
1182                    num_suffix: 0,
1183                    impl_related: true,
1184                }],
1185                conjunct_relations: vec![],
1186                primary_keys: vec![
1187                    PrimaryKey {
1188                        name: "cake_id".to_owned(),
1189                    },
1190                    PrimaryKey {
1191                        name: "filling_id".to_owned(),
1192                    },
1193                ],
1194            },
1195            Entity {
1196                table_name: "filling".to_owned(),
1197                columns: vec![
1198                    Column {
1199                        name: "id".to_owned(),
1200                        col_type: ColumnType::Integer,
1201                        auto_increment: true,
1202                        not_null: true,
1203                        unique: false,
1204                    },
1205                    Column {
1206                        name: "name".to_owned(),
1207                        col_type: ColumnType::String(StringLen::N(255)),
1208                        auto_increment: false,
1209                        not_null: true,
1210                        unique: false,
1211                    },
1212                ],
1213                relations: vec![],
1214                conjunct_relations: vec![ConjunctRelation {
1215                    via: "cake_filling".to_owned(),
1216                    to: "cake".to_owned(),
1217                }],
1218                primary_keys: vec![PrimaryKey {
1219                    name: "id".to_owned(),
1220                }],
1221            },
1222            Entity {
1223                table_name: "fruit".to_owned(),
1224                columns: vec![
1225                    Column {
1226                        name: "id".to_owned(),
1227                        col_type: ColumnType::Integer,
1228                        auto_increment: true,
1229                        not_null: true,
1230                        unique: false,
1231                    },
1232                    Column {
1233                        name: "name".to_owned(),
1234                        col_type: ColumnType::String(StringLen::N(255)),
1235                        auto_increment: false,
1236                        not_null: true,
1237                        unique: false,
1238                    },
1239                    Column {
1240                        name: "cake_id".to_owned(),
1241                        col_type: ColumnType::Integer,
1242                        auto_increment: false,
1243                        not_null: false,
1244                        unique: false,
1245                    },
1246                ],
1247                relations: vec![
1248                    Relation {
1249                        ref_table: "cake".to_owned(),
1250                        columns: vec!["cake_id".to_owned()],
1251                        ref_columns: vec!["id".to_owned()],
1252                        rel_type: RelationType::BelongsTo,
1253                        on_delete: None,
1254                        on_update: None,
1255                        self_referencing: false,
1256                        num_suffix: 0,
1257                        impl_related: true,
1258                    },
1259                    Relation {
1260                        ref_table: "vendor".to_owned(),
1261                        columns: vec![],
1262                        ref_columns: vec![],
1263                        rel_type: RelationType::HasMany,
1264                        on_delete: None,
1265                        on_update: None,
1266                        self_referencing: false,
1267                        num_suffix: 0,
1268                        impl_related: true,
1269                    },
1270                ],
1271                conjunct_relations: vec![],
1272                primary_keys: vec![PrimaryKey {
1273                    name: "id".to_owned(),
1274                }],
1275            },
1276            Entity {
1277                table_name: "vendor".to_owned(),
1278                columns: vec![
1279                    Column {
1280                        name: "id".to_owned(),
1281                        col_type: ColumnType::Integer,
1282                        auto_increment: true,
1283                        not_null: true,
1284                        unique: false,
1285                    },
1286                    Column {
1287                        name: "_name_".to_owned(),
1288                        col_type: ColumnType::String(StringLen::N(255)),
1289                        auto_increment: false,
1290                        not_null: true,
1291                        unique: false,
1292                    },
1293                    Column {
1294                        name: "fruitId".to_owned(),
1295                        col_type: ColumnType::Integer,
1296                        auto_increment: false,
1297                        not_null: false,
1298                        unique: false,
1299                    },
1300                ],
1301                relations: vec![Relation {
1302                    ref_table: "fruit".to_owned(),
1303                    columns: vec!["fruitId".to_owned()],
1304                    ref_columns: vec!["id".to_owned()],
1305                    rel_type: RelationType::BelongsTo,
1306                    on_delete: None,
1307                    on_update: None,
1308                    self_referencing: false,
1309                    num_suffix: 0,
1310                    impl_related: true,
1311                }],
1312                conjunct_relations: vec![],
1313                primary_keys: vec![PrimaryKey {
1314                    name: "id".to_owned(),
1315                }],
1316            },
1317            Entity {
1318                table_name: "rust_keyword".to_owned(),
1319                columns: vec![
1320                    Column {
1321                        name: "id".to_owned(),
1322                        col_type: ColumnType::Integer,
1323                        auto_increment: true,
1324                        not_null: true,
1325                        unique: false,
1326                    },
1327                    Column {
1328                        name: "testing".to_owned(),
1329                        col_type: ColumnType::TinyInteger,
1330                        auto_increment: false,
1331                        not_null: true,
1332                        unique: false,
1333                    },
1334                    Column {
1335                        name: "rust".to_owned(),
1336                        col_type: ColumnType::TinyUnsigned,
1337                        auto_increment: false,
1338                        not_null: true,
1339                        unique: false,
1340                    },
1341                    Column {
1342                        name: "keywords".to_owned(),
1343                        col_type: ColumnType::SmallInteger,
1344                        auto_increment: false,
1345                        not_null: true,
1346                        unique: false,
1347                    },
1348                    Column {
1349                        name: "type".to_owned(),
1350                        col_type: ColumnType::SmallUnsigned,
1351                        auto_increment: false,
1352                        not_null: true,
1353                        unique: false,
1354                    },
1355                    Column {
1356                        name: "typeof".to_owned(),
1357                        col_type: ColumnType::Integer,
1358                        auto_increment: false,
1359                        not_null: true,
1360                        unique: false,
1361                    },
1362                    Column {
1363                        name: "crate".to_owned(),
1364                        col_type: ColumnType::Unsigned,
1365                        auto_increment: false,
1366                        not_null: true,
1367                        unique: false,
1368                    },
1369                    Column {
1370                        name: "self".to_owned(),
1371                        col_type: ColumnType::BigInteger,
1372                        auto_increment: false,
1373                        not_null: true,
1374                        unique: false,
1375                    },
1376                    Column {
1377                        name: "self_id1".to_owned(),
1378                        col_type: ColumnType::BigUnsigned,
1379                        auto_increment: false,
1380                        not_null: true,
1381                        unique: false,
1382                    },
1383                    Column {
1384                        name: "self_id2".to_owned(),
1385                        col_type: ColumnType::Integer,
1386                        auto_increment: false,
1387                        not_null: true,
1388                        unique: false,
1389                    },
1390                    Column {
1391                        name: "fruit_id1".to_owned(),
1392                        col_type: ColumnType::Integer,
1393                        auto_increment: false,
1394                        not_null: true,
1395                        unique: false,
1396                    },
1397                    Column {
1398                        name: "fruit_id2".to_owned(),
1399                        col_type: ColumnType::Integer,
1400                        auto_increment: false,
1401                        not_null: true,
1402                        unique: false,
1403                    },
1404                    Column {
1405                        name: "cake_id".to_owned(),
1406                        col_type: ColumnType::Integer,
1407                        auto_increment: false,
1408                        not_null: true,
1409                        unique: false,
1410                    },
1411                ],
1412                relations: vec![
1413                    Relation {
1414                        ref_table: "rust_keyword".to_owned(),
1415                        columns: vec!["self_id1".to_owned()],
1416                        ref_columns: vec!["id".to_owned()],
1417                        rel_type: RelationType::BelongsTo,
1418                        on_delete: None,
1419                        on_update: None,
1420                        self_referencing: true,
1421                        num_suffix: 1,
1422                        impl_related: true,
1423                    },
1424                    Relation {
1425                        ref_table: "rust_keyword".to_owned(),
1426                        columns: vec!["self_id2".to_owned()],
1427                        ref_columns: vec!["id".to_owned()],
1428                        rel_type: RelationType::BelongsTo,
1429                        on_delete: None,
1430                        on_update: None,
1431                        self_referencing: true,
1432                        num_suffix: 2,
1433                        impl_related: true,
1434                    },
1435                    Relation {
1436                        ref_table: "fruit".to_owned(),
1437                        columns: vec!["fruit_id1".to_owned()],
1438                        ref_columns: vec!["id".to_owned()],
1439                        rel_type: RelationType::BelongsTo,
1440                        on_delete: None,
1441                        on_update: None,
1442                        self_referencing: false,
1443                        num_suffix: 1,
1444                        impl_related: true,
1445                    },
1446                    Relation {
1447                        ref_table: "fruit".to_owned(),
1448                        columns: vec!["fruit_id2".to_owned()],
1449                        ref_columns: vec!["id".to_owned()],
1450                        rel_type: RelationType::BelongsTo,
1451                        on_delete: None,
1452                        on_update: None,
1453                        self_referencing: false,
1454                        num_suffix: 2,
1455                        impl_related: true,
1456                    },
1457                    Relation {
1458                        ref_table: "cake".to_owned(),
1459                        columns: vec!["cake_id".to_owned()],
1460                        ref_columns: vec!["id".to_owned()],
1461                        rel_type: RelationType::BelongsTo,
1462                        on_delete: None,
1463                        on_update: None,
1464                        self_referencing: false,
1465                        num_suffix: 0,
1466                        impl_related: true,
1467                    },
1468                ],
1469                conjunct_relations: vec![],
1470                primary_keys: vec![PrimaryKey {
1471                    name: "id".to_owned(),
1472                }],
1473            },
1474            Entity {
1475                table_name: "cake_with_float".to_owned(),
1476                columns: vec![
1477                    Column {
1478                        name: "id".to_owned(),
1479                        col_type: ColumnType::Integer,
1480                        auto_increment: true,
1481                        not_null: true,
1482                        unique: false,
1483                    },
1484                    Column {
1485                        name: "name".to_owned(),
1486                        col_type: ColumnType::Text,
1487                        auto_increment: false,
1488                        not_null: false,
1489                        unique: false,
1490                    },
1491                    Column {
1492                        name: "price".to_owned(),
1493                        col_type: ColumnType::Float,
1494                        auto_increment: false,
1495                        not_null: false,
1496                        unique: false,
1497                    },
1498                ],
1499                relations: vec![Relation {
1500                    ref_table: "fruit".to_owned(),
1501                    columns: vec![],
1502                    ref_columns: vec![],
1503                    rel_type: RelationType::HasMany,
1504                    on_delete: None,
1505                    on_update: None,
1506                    self_referencing: false,
1507                    num_suffix: 0,
1508                    impl_related: true,
1509                }],
1510                conjunct_relations: vec![ConjunctRelation {
1511                    via: "cake_filling".to_owned(),
1512                    to: "filling".to_owned(),
1513                }],
1514                primary_keys: vec![PrimaryKey {
1515                    name: "id".to_owned(),
1516                }],
1517            },
1518            Entity {
1519                table_name: "cake_with_double".to_owned(),
1520                columns: vec![
1521                    Column {
1522                        name: "id".to_owned(),
1523                        col_type: ColumnType::Integer,
1524                        auto_increment: true,
1525                        not_null: true,
1526                        unique: false,
1527                    },
1528                    Column {
1529                        name: "name".to_owned(),
1530                        col_type: ColumnType::Text,
1531                        auto_increment: false,
1532                        not_null: false,
1533                        unique: false,
1534                    },
1535                    Column {
1536                        name: "price".to_owned(),
1537                        col_type: ColumnType::Double,
1538                        auto_increment: false,
1539                        not_null: false,
1540                        unique: false,
1541                    },
1542                ],
1543                relations: vec![Relation {
1544                    ref_table: "fruit".to_owned(),
1545                    columns: vec![],
1546                    ref_columns: vec![],
1547                    rel_type: RelationType::HasMany,
1548                    on_delete: None,
1549                    on_update: None,
1550                    self_referencing: false,
1551                    num_suffix: 0,
1552                    impl_related: true,
1553                }],
1554                conjunct_relations: vec![ConjunctRelation {
1555                    via: "cake_filling".to_owned(),
1556                    to: "filling".to_owned(),
1557                }],
1558                primary_keys: vec![PrimaryKey {
1559                    name: "id".to_owned(),
1560                }],
1561            },
1562            Entity {
1563                table_name: "collection".to_owned(),
1564                columns: vec![
1565                    Column {
1566                        name: "id".to_owned(),
1567                        col_type: ColumnType::Integer,
1568                        auto_increment: true,
1569                        not_null: true,
1570                        unique: false,
1571                    },
1572                    Column {
1573                        name: "integers".to_owned(),
1574                        col_type: ColumnType::Array(RcOrArc::new(ColumnType::Integer)),
1575                        auto_increment: false,
1576                        not_null: true,
1577                        unique: false,
1578                    },
1579                    Column {
1580                        name: "integers_opt".to_owned(),
1581                        col_type: ColumnType::Array(RcOrArc::new(ColumnType::Integer)),
1582                        auto_increment: false,
1583                        not_null: false,
1584                        unique: false,
1585                    },
1586                ],
1587                relations: vec![],
1588                conjunct_relations: vec![],
1589                primary_keys: vec![PrimaryKey {
1590                    name: "id".to_owned(),
1591                }],
1592            },
1593            Entity {
1594                table_name: "collection_float".to_owned(),
1595                columns: vec![
1596                    Column {
1597                        name: "id".to_owned(),
1598                        col_type: ColumnType::Integer,
1599                        auto_increment: true,
1600                        not_null: true,
1601                        unique: false,
1602                    },
1603                    Column {
1604                        name: "floats".to_owned(),
1605                        col_type: ColumnType::Array(RcOrArc::new(ColumnType::Float)),
1606                        auto_increment: false,
1607                        not_null: true,
1608                        unique: false,
1609                    },
1610                    Column {
1611                        name: "doubles".to_owned(),
1612                        col_type: ColumnType::Array(RcOrArc::new(ColumnType::Double)),
1613                        auto_increment: false,
1614                        not_null: true,
1615                        unique: false,
1616                    },
1617                ],
1618                relations: vec![],
1619                conjunct_relations: vec![],
1620                primary_keys: vec![PrimaryKey {
1621                    name: "id".to_owned(),
1622                }],
1623            },
1624            Entity {
1625                table_name: "parent".to_owned(),
1626                columns: vec![
1627                    Column {
1628                        name: "id1".to_owned(),
1629                        col_type: ColumnType::Integer,
1630                        auto_increment: false,
1631                        not_null: true,
1632                        unique: false,
1633                    },
1634                    Column {
1635                        name: "id2".to_owned(),
1636                        col_type: ColumnType::Integer,
1637                        auto_increment: false,
1638                        not_null: true,
1639                        unique: false,
1640                    },
1641                ],
1642                relations: vec![Relation {
1643                    ref_table: "child".to_owned(),
1644                    columns: vec![],
1645                    ref_columns: vec![],
1646                    rel_type: RelationType::HasMany,
1647                    on_delete: None,
1648                    on_update: None,
1649                    self_referencing: false,
1650                    num_suffix: 0,
1651                    impl_related: true,
1652                }],
1653                conjunct_relations: vec![],
1654                primary_keys: vec![
1655                    PrimaryKey {
1656                        name: "id1".to_owned(),
1657                    },
1658                    PrimaryKey {
1659                        name: "id2".to_owned(),
1660                    },
1661                ],
1662            },
1663            Entity {
1664                table_name: "child".to_owned(),
1665                columns: vec![
1666                    Column {
1667                        name: "id".to_owned(),
1668                        col_type: ColumnType::Integer,
1669                        auto_increment: true,
1670                        not_null: true,
1671                        unique: false,
1672                    },
1673                    Column {
1674                        name: "parent_id1".to_owned(),
1675                        col_type: ColumnType::Integer,
1676                        auto_increment: false,
1677                        not_null: true,
1678                        unique: false,
1679                    },
1680                    Column {
1681                        name: "parent_id2".to_owned(),
1682                        col_type: ColumnType::Integer,
1683                        auto_increment: false,
1684                        not_null: true,
1685                        unique: false,
1686                    },
1687                ],
1688                relations: vec![Relation {
1689                    ref_table: "parent".to_owned(),
1690                    columns: vec!["parent_id1".to_owned(), "parent_id2".to_owned()],
1691                    ref_columns: vec!["id1".to_owned(), "id2".to_owned()],
1692                    rel_type: RelationType::BelongsTo,
1693                    on_delete: None,
1694                    on_update: None,
1695                    self_referencing: false,
1696                    num_suffix: 0,
1697                    impl_related: true,
1698                }],
1699                conjunct_relations: vec![],
1700                primary_keys: vec![PrimaryKey {
1701                    name: "id".to_owned(),
1702                }],
1703            },
1704        ]
1705    }
1706
1707    fn parse_from_file<R>(inner: R) -> io::Result<TokenStream>
1708    where
1709        R: Read,
1710    {
1711        let mut reader = BufReader::new(inner);
1712        let mut lines: Vec<String> = Vec::new();
1713
1714        reader.read_until(b';', &mut Vec::new())?;
1715
1716        let mut line = String::new();
1717        while reader.read_line(&mut line)? > 0 {
1718            lines.push(line.to_owned());
1719            line.clear();
1720        }
1721        let content = lines.join("");
1722        Ok(content.parse().unwrap())
1723    }
1724
1725    fn parse_from_frontend_file<R>(inner: R) -> io::Result<TokenStream>
1726    where
1727        R: Read,
1728    {
1729        let mut reader = BufReader::new(inner);
1730        let mut lines: Vec<String> = Vec::new();
1731
1732        reader.read_until(b'\n', &mut Vec::new())?;
1733
1734        let mut line = String::new();
1735        while reader.read_line(&mut line)? > 0 {
1736            lines.push(line.to_owned());
1737            line.clear();
1738        }
1739        let content = lines.join("");
1740        Ok(content.parse().unwrap())
1741    }
1742
1743    #[test]
1744    fn test_gen_expanded_code_blocks() -> io::Result<()> {
1745        let entities = setup();
1746        const ENTITY_FILES: [&str; 13] = [
1747            include_str!("../../tests/expanded/cake.rs"),
1748            include_str!("../../tests/expanded/cake_filling.rs"),
1749            include_str!("../../tests/expanded/cake_filling_price.rs"),
1750            include_str!("../../tests/expanded/filling.rs"),
1751            include_str!("../../tests/expanded/fruit.rs"),
1752            include_str!("../../tests/expanded/vendor.rs"),
1753            include_str!("../../tests/expanded/rust_keyword.rs"),
1754            include_str!("../../tests/expanded/cake_with_float.rs"),
1755            include_str!("../../tests/expanded/cake_with_double.rs"),
1756            include_str!("../../tests/expanded/collection.rs"),
1757            include_str!("../../tests/expanded/collection_float.rs"),
1758            include_str!("../../tests/expanded/parent.rs"),
1759            include_str!("../../tests/expanded/child.rs"),
1760        ];
1761        const ENTITY_FILES_WITH_SCHEMA_NAME: [&str; 13] = [
1762            include_str!("../../tests/expanded_with_schema_name/cake.rs"),
1763            include_str!("../../tests/expanded_with_schema_name/cake_filling.rs"),
1764            include_str!("../../tests/expanded_with_schema_name/cake_filling_price.rs"),
1765            include_str!("../../tests/expanded_with_schema_name/filling.rs"),
1766            include_str!("../../tests/expanded_with_schema_name/fruit.rs"),
1767            include_str!("../../tests/expanded_with_schema_name/vendor.rs"),
1768            include_str!("../../tests/expanded_with_schema_name/rust_keyword.rs"),
1769            include_str!("../../tests/expanded_with_schema_name/cake_with_float.rs"),
1770            include_str!("../../tests/expanded_with_schema_name/cake_with_double.rs"),
1771            include_str!("../../tests/expanded_with_schema_name/collection.rs"),
1772            include_str!("../../tests/expanded_with_schema_name/collection_float.rs"),
1773            include_str!("../../tests/expanded_with_schema_name/parent.rs"),
1774            include_str!("../../tests/expanded_with_schema_name/child.rs"),
1775        ];
1776
1777        assert_eq!(entities.len(), ENTITY_FILES.len());
1778
1779        for (i, entity) in entities.iter().enumerate() {
1780            assert_eq!(
1781                parse_from_file(ENTITY_FILES[i].as_bytes())?.to_string(),
1782                EntityWriter::gen_expanded_code_blocks(
1783                    entity,
1784                    &crate::WithSerde::None,
1785                    &crate::DateTimeCrate::Chrono,
1786                    &None,
1787                    false,
1788                    false,
1789                    &TokenStream::new(),
1790                    &TokenStream::new(),
1791                    false,
1792                    true,
1793                )
1794                .into_iter()
1795                .skip(1)
1796                .fold(TokenStream::new(), |mut acc, tok| {
1797                    acc.extend(tok);
1798                    acc
1799                })
1800                .to_string()
1801            );
1802            assert_eq!(
1803                parse_from_file(ENTITY_FILES_WITH_SCHEMA_NAME[i].as_bytes())?.to_string(),
1804                EntityWriter::gen_expanded_code_blocks(
1805                    entity,
1806                    &crate::WithSerde::None,
1807                    &crate::DateTimeCrate::Chrono,
1808                    &Some("schema_name".to_owned()),
1809                    false,
1810                    false,
1811                    &TokenStream::new(),
1812                    &TokenStream::new(),
1813                    false,
1814                    true,
1815                )
1816                .into_iter()
1817                .skip(1)
1818                .fold(TokenStream::new(), |mut acc, tok| {
1819                    acc.extend(tok);
1820                    acc
1821                })
1822                .to_string()
1823            );
1824        }
1825
1826        Ok(())
1827    }
1828
1829    #[test]
1830    fn test_gen_compact_code_blocks() -> io::Result<()> {
1831        let entities = setup();
1832        const ENTITY_FILES: [&str; 13] = [
1833            include_str!("../../tests/compact/cake.rs"),
1834            include_str!("../../tests/compact/cake_filling.rs"),
1835            include_str!("../../tests/compact/cake_filling_price.rs"),
1836            include_str!("../../tests/compact/filling.rs"),
1837            include_str!("../../tests/compact/fruit.rs"),
1838            include_str!("../../tests/compact/vendor.rs"),
1839            include_str!("../../tests/compact/rust_keyword.rs"),
1840            include_str!("../../tests/compact/cake_with_float.rs"),
1841            include_str!("../../tests/compact/cake_with_double.rs"),
1842            include_str!("../../tests/compact/collection.rs"),
1843            include_str!("../../tests/compact/collection_float.rs"),
1844            include_str!("../../tests/compact/parent.rs"),
1845            include_str!("../../tests/compact/child.rs"),
1846        ];
1847        const ENTITY_FILES_WITH_SCHEMA_NAME: [&str; 13] = [
1848            include_str!("../../tests/compact_with_schema_name/cake.rs"),
1849            include_str!("../../tests/compact_with_schema_name/cake_filling.rs"),
1850            include_str!("../../tests/compact_with_schema_name/cake_filling_price.rs"),
1851            include_str!("../../tests/compact_with_schema_name/filling.rs"),
1852            include_str!("../../tests/compact_with_schema_name/fruit.rs"),
1853            include_str!("../../tests/compact_with_schema_name/vendor.rs"),
1854            include_str!("../../tests/compact_with_schema_name/rust_keyword.rs"),
1855            include_str!("../../tests/compact_with_schema_name/cake_with_float.rs"),
1856            include_str!("../../tests/compact_with_schema_name/cake_with_double.rs"),
1857            include_str!("../../tests/compact_with_schema_name/collection.rs"),
1858            include_str!("../../tests/compact_with_schema_name/collection_float.rs"),
1859            include_str!("../../tests/compact_with_schema_name/parent.rs"),
1860            include_str!("../../tests/compact_with_schema_name/child.rs"),
1861        ];
1862
1863        assert_eq!(entities.len(), ENTITY_FILES.len());
1864
1865        for (i, entity) in entities.iter().enumerate() {
1866            assert_eq!(
1867                parse_from_file(ENTITY_FILES[i].as_bytes())?.to_string(),
1868                EntityWriter::gen_compact_code_blocks(
1869                    entity,
1870                    &crate::WithSerde::None,
1871                    &crate::DateTimeCrate::Chrono,
1872                    &None,
1873                    false,
1874                    false,
1875                    &TokenStream::new(),
1876                    &TokenStream::new(),
1877                    false,
1878                    true,
1879                )
1880                .into_iter()
1881                .skip(1)
1882                .fold(TokenStream::new(), |mut acc, tok| {
1883                    acc.extend(tok);
1884                    acc
1885                })
1886                .to_string()
1887            );
1888            assert_eq!(
1889                parse_from_file(ENTITY_FILES_WITH_SCHEMA_NAME[i].as_bytes())?.to_string(),
1890                EntityWriter::gen_compact_code_blocks(
1891                    entity,
1892                    &crate::WithSerde::None,
1893                    &crate::DateTimeCrate::Chrono,
1894                    &Some("schema_name".to_owned()),
1895                    false,
1896                    false,
1897                    &TokenStream::new(),
1898                    &TokenStream::new(),
1899                    false,
1900                    true,
1901                )
1902                .into_iter()
1903                .skip(1)
1904                .fold(TokenStream::new(), |mut acc, tok| {
1905                    acc.extend(tok);
1906                    acc
1907                })
1908                .to_string()
1909            );
1910        }
1911
1912        Ok(())
1913    }
1914
1915    #[test]
1916    fn test_gen_frontend_code_blocks() -> io::Result<()> {
1917        let entities = setup();
1918        const ENTITY_FILES: [&str; 13] = [
1919            include_str!("../../tests/frontend/cake.rs"),
1920            include_str!("../../tests/frontend/cake_filling.rs"),
1921            include_str!("../../tests/frontend/cake_filling_price.rs"),
1922            include_str!("../../tests/frontend/filling.rs"),
1923            include_str!("../../tests/frontend/fruit.rs"),
1924            include_str!("../../tests/frontend/vendor.rs"),
1925            include_str!("../../tests/frontend/rust_keyword.rs"),
1926            include_str!("../../tests/frontend/cake_with_float.rs"),
1927            include_str!("../../tests/frontend/cake_with_double.rs"),
1928            include_str!("../../tests/frontend/collection.rs"),
1929            include_str!("../../tests/frontend/collection_float.rs"),
1930            include_str!("../../tests/frontend/parent.rs"),
1931            include_str!("../../tests/frontend/child.rs"),
1932        ];
1933        const ENTITY_FILES_WITH_SCHEMA_NAME: [&str; 13] = [
1934            include_str!("../../tests/frontend_with_schema_name/cake.rs"),
1935            include_str!("../../tests/frontend_with_schema_name/cake_filling.rs"),
1936            include_str!("../../tests/frontend_with_schema_name/cake_filling_price.rs"),
1937            include_str!("../../tests/frontend_with_schema_name/filling.rs"),
1938            include_str!("../../tests/frontend_with_schema_name/fruit.rs"),
1939            include_str!("../../tests/frontend_with_schema_name/vendor.rs"),
1940            include_str!("../../tests/frontend_with_schema_name/rust_keyword.rs"),
1941            include_str!("../../tests/frontend_with_schema_name/cake_with_float.rs"),
1942            include_str!("../../tests/frontend_with_schema_name/cake_with_double.rs"),
1943            include_str!("../../tests/frontend_with_schema_name/collection.rs"),
1944            include_str!("../../tests/frontend_with_schema_name/collection_float.rs"),
1945            include_str!("../../tests/frontend_with_schema_name/parent.rs"),
1946            include_str!("../../tests/frontend_with_schema_name/child.rs"),
1947        ];
1948
1949        assert_eq!(entities.len(), ENTITY_FILES.len());
1950
1951        for (i, entity) in entities.iter().enumerate() {
1952            assert_eq!(
1953                dbg!(parse_from_frontend_file(ENTITY_FILES[i].as_bytes())?.to_string()),
1954                EntityWriter::gen_frontend_code_blocks(
1955                    entity,
1956                    &crate::WithSerde::None,
1957                    &crate::DateTimeCrate::Chrono,
1958                    &None,
1959                    false,
1960                    false,
1961                    &TokenStream::new(),
1962                    &TokenStream::new(),
1963                    false,
1964                    true,
1965                )
1966                .into_iter()
1967                .skip(1)
1968                .fold(TokenStream::new(), |mut acc, tok| {
1969                    acc.extend(tok);
1970                    acc
1971                })
1972                .to_string()
1973            );
1974            assert_eq!(
1975                parse_from_frontend_file(ENTITY_FILES_WITH_SCHEMA_NAME[i].as_bytes())?.to_string(),
1976                EntityWriter::gen_frontend_code_blocks(
1977                    entity,
1978                    &crate::WithSerde::None,
1979                    &crate::DateTimeCrate::Chrono,
1980                    &Some("schema_name".to_owned()),
1981                    false,
1982                    false,
1983                    &TokenStream::new(),
1984                    &TokenStream::new(),
1985                    false,
1986                    true,
1987                )
1988                .into_iter()
1989                .skip(1)
1990                .fold(TokenStream::new(), |mut acc, tok| {
1991                    acc.extend(tok);
1992                    acc
1993                })
1994                .to_string()
1995            );
1996        }
1997
1998        Ok(())
1999    }
2000
2001    #[test]
2002    fn test_gen_with_serde() -> io::Result<()> {
2003        let cake_entity = setup().get(0).unwrap().clone();
2004
2005        assert_eq!(cake_entity.get_table_name_snake_case(), "cake");
2006
2007        // Compact code blocks
2008        assert_eq!(
2009            comparable_file_string(include_str!("../../tests/compact_with_serde/cake_none.rs"))?,
2010            generated_to_string(EntityWriter::gen_compact_code_blocks(
2011                &cake_entity,
2012                &WithSerde::None,
2013                &DateTimeCrate::Chrono,
2014                &None,
2015                false,
2016                false,
2017                &TokenStream::new(),
2018                &TokenStream::new(),
2019                false,
2020                true,
2021            ))
2022        );
2023        assert_eq!(
2024            comparable_file_string(include_str!(
2025                "../../tests/compact_with_serde/cake_serialize.rs"
2026            ))?,
2027            generated_to_string(EntityWriter::gen_compact_code_blocks(
2028                &cake_entity,
2029                &WithSerde::Serialize,
2030                &DateTimeCrate::Chrono,
2031                &None,
2032                false,
2033                false,
2034                &TokenStream::new(),
2035                &TokenStream::new(),
2036                false,
2037                true,
2038            ))
2039        );
2040        assert_eq!(
2041            comparable_file_string(include_str!(
2042                "../../tests/compact_with_serde/cake_deserialize.rs"
2043            ))?,
2044            generated_to_string(EntityWriter::gen_compact_code_blocks(
2045                &cake_entity,
2046                &WithSerde::Deserialize,
2047                &DateTimeCrate::Chrono,
2048                &None,
2049                true,
2050                false,
2051                &TokenStream::new(),
2052                &TokenStream::new(),
2053                false,
2054                true,
2055            ))
2056        );
2057        assert_eq!(
2058            comparable_file_string(include_str!("../../tests/compact_with_serde/cake_both.rs"))?,
2059            generated_to_string(EntityWriter::gen_compact_code_blocks(
2060                &cake_entity,
2061                &WithSerde::Both,
2062                &DateTimeCrate::Chrono,
2063                &None,
2064                true,
2065                false,
2066                &TokenStream::new(),
2067                &TokenStream::new(),
2068                false,
2069                true,
2070            ))
2071        );
2072
2073        // Expanded code blocks
2074        assert_eq!(
2075            comparable_file_string(include_str!("../../tests/expanded_with_serde/cake_none.rs"))?,
2076            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2077                &cake_entity,
2078                &WithSerde::None,
2079                &DateTimeCrate::Chrono,
2080                &None,
2081                false,
2082                false,
2083                &TokenStream::new(),
2084                &TokenStream::new(),
2085                false,
2086                true,
2087            ))
2088        );
2089        assert_eq!(
2090            comparable_file_string(include_str!(
2091                "../../tests/expanded_with_serde/cake_serialize.rs"
2092            ))?,
2093            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2094                &cake_entity,
2095                &WithSerde::Serialize,
2096                &DateTimeCrate::Chrono,
2097                &None,
2098                false,
2099                false,
2100                &TokenStream::new(),
2101                &TokenStream::new(),
2102                false,
2103                true,
2104            ))
2105        );
2106        assert_eq!(
2107            comparable_file_string(include_str!(
2108                "../../tests/expanded_with_serde/cake_deserialize.rs"
2109            ))?,
2110            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2111                &cake_entity,
2112                &WithSerde::Deserialize,
2113                &DateTimeCrate::Chrono,
2114                &None,
2115                true,
2116                false,
2117                &TokenStream::new(),
2118                &TokenStream::new(),
2119                false,
2120                true,
2121            ))
2122        );
2123        assert_eq!(
2124            comparable_file_string(include_str!("../../tests/expanded_with_serde/cake_both.rs"))?,
2125            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2126                &cake_entity,
2127                &WithSerde::Both,
2128                &DateTimeCrate::Chrono,
2129                &None,
2130                true,
2131                false,
2132                &TokenStream::new(),
2133                &TokenStream::new(),
2134                false,
2135                true,
2136            ))
2137        );
2138
2139        // Frontend code blocks
2140        assert_eq!(
2141            comparable_file_string(include_str!("../../tests/frontend_with_serde/cake_none.rs"))?,
2142            generated_to_string(EntityWriter::gen_frontend_code_blocks(
2143                &cake_entity,
2144                &WithSerde::None,
2145                &DateTimeCrate::Chrono,
2146                &None,
2147                false,
2148                false,
2149                &TokenStream::new(),
2150                &TokenStream::new(),
2151                false,
2152                true,
2153            ))
2154        );
2155        assert_eq!(
2156            comparable_file_string(include_str!(
2157                "../../tests/frontend_with_serde/cake_serialize.rs"
2158            ))?,
2159            generated_to_string(EntityWriter::gen_frontend_code_blocks(
2160                &cake_entity,
2161                &WithSerde::Serialize,
2162                &DateTimeCrate::Chrono,
2163                &None,
2164                false,
2165                false,
2166                &TokenStream::new(),
2167                &TokenStream::new(),
2168                false,
2169                true,
2170            ))
2171        );
2172        assert_eq!(
2173            comparable_file_string(include_str!(
2174                "../../tests/frontend_with_serde/cake_deserialize.rs"
2175            ))?,
2176            generated_to_string(EntityWriter::gen_frontend_code_blocks(
2177                &cake_entity,
2178                &WithSerde::Deserialize,
2179                &DateTimeCrate::Chrono,
2180                &None,
2181                true,
2182                false,
2183                &TokenStream::new(),
2184                &TokenStream::new(),
2185                false,
2186                true,
2187            ))
2188        );
2189        assert_eq!(
2190            comparable_file_string(include_str!("../../tests/frontend_with_serde/cake_both.rs"))?,
2191            generated_to_string(EntityWriter::gen_frontend_code_blocks(
2192                &cake_entity,
2193                &WithSerde::Both,
2194                &DateTimeCrate::Chrono,
2195                &None,
2196                true,
2197                false,
2198                &TokenStream::new(),
2199                &TokenStream::new(),
2200                false,
2201                true,
2202            ))
2203        );
2204
2205        Ok(())
2206    }
2207
2208    #[test]
2209    fn test_gen_with_seaography() -> io::Result<()> {
2210        let cake_entity = Entity {
2211            table_name: "cake".to_owned(),
2212            columns: vec![
2213                Column {
2214                    name: "id".to_owned(),
2215                    col_type: ColumnType::Integer,
2216                    auto_increment: true,
2217                    not_null: true,
2218                    unique: false,
2219                },
2220                Column {
2221                    name: "name".to_owned(),
2222                    col_type: ColumnType::Text,
2223                    auto_increment: false,
2224                    not_null: false,
2225                    unique: false,
2226                },
2227                Column {
2228                    name: "base_id".to_owned(),
2229                    col_type: ColumnType::Integer,
2230                    auto_increment: false,
2231                    not_null: false,
2232                    unique: false,
2233                },
2234            ],
2235            relations: vec![
2236                Relation {
2237                    ref_table: "fruit".to_owned(),
2238                    columns: vec![],
2239                    ref_columns: vec![],
2240                    rel_type: RelationType::HasMany,
2241                    on_delete: None,
2242                    on_update: None,
2243                    self_referencing: false,
2244                    num_suffix: 0,
2245                    impl_related: true,
2246                },
2247                Relation {
2248                    ref_table: "cake".to_owned(),
2249                    columns: vec![],
2250                    ref_columns: vec![],
2251                    rel_type: RelationType::HasOne,
2252                    on_delete: None,
2253                    on_update: None,
2254                    self_referencing: true,
2255                    num_suffix: 0,
2256                    impl_related: true,
2257                },
2258            ],
2259            conjunct_relations: vec![ConjunctRelation {
2260                via: "cake_filling".to_owned(),
2261                to: "filling".to_owned(),
2262            }],
2263            primary_keys: vec![PrimaryKey {
2264                name: "id".to_owned(),
2265            }],
2266        };
2267
2268        assert_eq!(cake_entity.get_table_name_snake_case(), "cake");
2269
2270        // Compact code blocks
2271        assert_eq!(
2272            comparable_file_string(include_str!("../../tests/with_seaography/cake.rs"))?,
2273            generated_to_string(EntityWriter::gen_compact_code_blocks(
2274                &cake_entity,
2275                &WithSerde::None,
2276                &DateTimeCrate::Chrono,
2277                &None,
2278                false,
2279                false,
2280                &TokenStream::new(),
2281                &TokenStream::new(),
2282                true,
2283                true,
2284            ))
2285        );
2286
2287        // Expanded code blocks
2288        assert_eq!(
2289            comparable_file_string(include_str!("../../tests/with_seaography/cake_expanded.rs"))?,
2290            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2291                &cake_entity,
2292                &WithSerde::None,
2293                &DateTimeCrate::Chrono,
2294                &None,
2295                false,
2296                false,
2297                &TokenStream::new(),
2298                &TokenStream::new(),
2299                true,
2300                true,
2301            ))
2302        );
2303
2304        // Frontend code blocks
2305        assert_eq!(
2306            comparable_file_string(include_str!("../../tests/with_seaography/cake_frontend.rs"))?,
2307            generated_to_string(EntityWriter::gen_frontend_code_blocks(
2308                &cake_entity,
2309                &WithSerde::None,
2310                &DateTimeCrate::Chrono,
2311                &None,
2312                false,
2313                false,
2314                &TokenStream::new(),
2315                &TokenStream::new(),
2316                true,
2317                true,
2318            ))
2319        );
2320
2321        Ok(())
2322    }
2323
2324    #[test]
2325    fn test_gen_with_seaography_mod() -> io::Result<()> {
2326        use crate::ActiveEnum;
2327        use sea_query::IntoIden;
2328
2329        let entities = setup();
2330        let enums = vec![
2331            (
2332                "coinflip_result_type",
2333                ActiveEnum {
2334                    enum_name: Alias::new("coinflip_result_type").into_iden(),
2335                    values: vec!["HEADS", "TAILS"]
2336                        .into_iter()
2337                        .map(|variant| Alias::new(variant).into_iden())
2338                        .collect(),
2339                },
2340            ),
2341            (
2342                "media_type",
2343                ActiveEnum {
2344                    enum_name: Alias::new("media_type").into_iden(),
2345                    values: vec![
2346                        "UNKNOWN",
2347                        "BITMAP",
2348                        "DRAWING",
2349                        "AUDIO",
2350                        "VIDEO",
2351                        "MULTIMEDIA",
2352                        "OFFICE",
2353                        "TEXT",
2354                        "EXECUTABLE",
2355                        "ARCHIVE",
2356                        "3D",
2357                    ]
2358                    .into_iter()
2359                    .map(|variant| Alias::new(variant).into_iden())
2360                    .collect(),
2361                },
2362            ),
2363        ]
2364        .into_iter()
2365        .map(|(k, v)| (k.to_string(), v))
2366        .collect();
2367
2368        assert_eq!(
2369            comparable_file_string(include_str!("../../tests/with_seaography/mod.rs"))?,
2370            generated_to_string(vec![EntityWriter::gen_seaography_entity_mod(
2371                &entities, &enums,
2372            )])
2373        );
2374
2375        Ok(())
2376    }
2377
2378    #[test]
2379    fn test_gen_with_derives() -> io::Result<()> {
2380        let mut cake_entity = setup().get_mut(0).unwrap().clone();
2381
2382        assert_eq!(cake_entity.get_table_name_snake_case(), "cake");
2383
2384        // Compact code blocks
2385        assert_eq!(
2386            comparable_file_string(include_str!(
2387                "../../tests/compact_with_derives/cake_none.rs"
2388            ))?,
2389            generated_to_string(EntityWriter::gen_compact_code_blocks(
2390                &cake_entity,
2391                &WithSerde::None,
2392                &DateTimeCrate::Chrono,
2393                &None,
2394                false,
2395                false,
2396                &TokenStream::new(),
2397                &TokenStream::new(),
2398                false,
2399                true,
2400            ))
2401        );
2402        assert_eq!(
2403            comparable_file_string(include_str!("../../tests/compact_with_derives/cake_one.rs"))?,
2404            generated_to_string(EntityWriter::gen_compact_code_blocks(
2405                &cake_entity,
2406                &WithSerde::None,
2407                &DateTimeCrate::Chrono,
2408                &None,
2409                false,
2410                false,
2411                &bonus_derive(["ts_rs::TS"]),
2412                &TokenStream::new(),
2413                false,
2414                true,
2415            ))
2416        );
2417        assert_eq!(
2418            comparable_file_string(include_str!(
2419                "../../tests/compact_with_derives/cake_multiple.rs"
2420            ))?,
2421            generated_to_string(EntityWriter::gen_compact_code_blocks(
2422                &cake_entity,
2423                &WithSerde::None,
2424                &DateTimeCrate::Chrono,
2425                &None,
2426                false,
2427                false,
2428                &bonus_derive(["ts_rs::TS", "utoipa::ToSchema"]),
2429                &TokenStream::new(),
2430                false,
2431                true,
2432            ))
2433        );
2434
2435        // Expanded code blocks
2436        assert_eq!(
2437            comparable_file_string(include_str!(
2438                "../../tests/expanded_with_derives/cake_none.rs"
2439            ))?,
2440            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2441                &cake_entity,
2442                &WithSerde::None,
2443                &DateTimeCrate::Chrono,
2444                &None,
2445                false,
2446                false,
2447                &TokenStream::new(),
2448                &TokenStream::new(),
2449                false,
2450                true,
2451            ))
2452        );
2453        assert_eq!(
2454            comparable_file_string(include_str!(
2455                "../../tests/expanded_with_derives/cake_one.rs"
2456            ))?,
2457            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2458                &cake_entity,
2459                &WithSerde::None,
2460                &DateTimeCrate::Chrono,
2461                &None,
2462                false,
2463                false,
2464                &bonus_derive(["ts_rs::TS"]),
2465                &TokenStream::new(),
2466                false,
2467                true,
2468            ))
2469        );
2470        assert_eq!(
2471            comparable_file_string(include_str!(
2472                "../../tests/expanded_with_derives/cake_multiple.rs"
2473            ))?,
2474            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2475                &cake_entity,
2476                &WithSerde::None,
2477                &DateTimeCrate::Chrono,
2478                &None,
2479                false,
2480                false,
2481                &bonus_derive(["ts_rs::TS", "utoipa::ToSchema"]),
2482                &TokenStream::new(),
2483                false,
2484                true,
2485            ))
2486        );
2487
2488        // Frontend code blocks
2489        assert_eq!(
2490            comparable_file_string(include_str!(
2491                "../../tests/frontend_with_derives/cake_none.rs"
2492            ))?,
2493            generated_to_string(EntityWriter::gen_frontend_code_blocks(
2494                &cake_entity,
2495                &WithSerde::None,
2496                &DateTimeCrate::Chrono,
2497                &None,
2498                false,
2499                false,
2500                &TokenStream::new(),
2501                &TokenStream::new(),
2502                false,
2503                true,
2504            ))
2505        );
2506        assert_eq!(
2507            comparable_file_string(include_str!(
2508                "../../tests/frontend_with_derives/cake_one.rs"
2509            ))?,
2510            generated_to_string(EntityWriter::gen_frontend_code_blocks(
2511                &cake_entity,
2512                &WithSerde::None,
2513                &DateTimeCrate::Chrono,
2514                &None,
2515                false,
2516                false,
2517                &bonus_derive(["ts_rs::TS"]),
2518                &TokenStream::new(),
2519                false,
2520                true,
2521            ))
2522        );
2523        assert_eq!(
2524            comparable_file_string(include_str!(
2525                "../../tests/frontend_with_derives/cake_multiple.rs"
2526            ))?,
2527            generated_to_string(EntityWriter::gen_frontend_code_blocks(
2528                &cake_entity,
2529                &WithSerde::None,
2530                &DateTimeCrate::Chrono,
2531                &None,
2532                false,
2533                false,
2534                &bonus_derive(["ts_rs::TS", "utoipa::ToSchema"]),
2535                &TokenStream::new(),
2536                false,
2537                true,
2538            ))
2539        );
2540
2541        // Make the `name` column of `cake` entity as hidden column
2542        cake_entity.columns[1].name = "_name".into();
2543
2544        assert_serde_variant_results(
2545            &cake_entity,
2546            &(
2547                include_str!("../../tests/compact_with_serde/cake_serialize_with_hidden_column.rs"),
2548                WithSerde::Serialize,
2549                None,
2550            ),
2551            Box::new(EntityWriter::gen_compact_code_blocks),
2552        )?;
2553        assert_serde_variant_results(
2554            &cake_entity,
2555            &(
2556                include_str!(
2557                    "../../tests/expanded_with_serde/cake_serialize_with_hidden_column.rs"
2558                ),
2559                WithSerde::Serialize,
2560                None,
2561            ),
2562            Box::new(EntityWriter::gen_expanded_code_blocks),
2563        )?;
2564        assert_serde_variant_results(
2565            &cake_entity,
2566            &(
2567                include_str!(
2568                    "../../tests/frontend_with_serde/cake_serialize_with_hidden_column.rs"
2569                ),
2570                WithSerde::Serialize,
2571                None,
2572            ),
2573            Box::new(EntityWriter::gen_frontend_code_blocks),
2574        )?;
2575
2576        Ok(())
2577    }
2578
2579    #[allow(clippy::type_complexity)]
2580    fn assert_serde_variant_results(
2581        cake_entity: &Entity,
2582        entity_serde_variant: &(&str, WithSerde, Option<String>),
2583        generator: Box<
2584            dyn Fn(
2585                &Entity,
2586                &WithSerde,
2587                &DateTimeCrate,
2588                &Option<String>,
2589                bool,
2590                bool,
2591                &TokenStream,
2592                &TokenStream,
2593                bool,
2594                bool,
2595            ) -> Vec<TokenStream>,
2596        >,
2597    ) -> io::Result<()> {
2598        let mut reader = BufReader::new(entity_serde_variant.0.as_bytes());
2599        let mut lines: Vec<String> = Vec::new();
2600        let serde_skip_deserializing_primary_key = matches!(
2601            entity_serde_variant.1,
2602            WithSerde::Both | WithSerde::Deserialize
2603        );
2604        let serde_skip_hidden_column = matches!(entity_serde_variant.1, WithSerde::Serialize);
2605
2606        reader.read_until(b'\n', &mut Vec::new())?;
2607
2608        let mut line = String::new();
2609        while reader.read_line(&mut line)? > 0 {
2610            lines.push(line.to_owned());
2611            line.clear();
2612        }
2613        let content = lines.join("");
2614        let expected: TokenStream = content.parse().unwrap();
2615        println!("{:?}", entity_serde_variant.1);
2616        let generated = generator(
2617            cake_entity,
2618            &entity_serde_variant.1,
2619            &DateTimeCrate::Chrono,
2620            &entity_serde_variant.2,
2621            serde_skip_deserializing_primary_key,
2622            serde_skip_hidden_column,
2623            &TokenStream::new(),
2624            &TokenStream::new(),
2625            false,
2626            true,
2627        )
2628        .into_iter()
2629        .fold(TokenStream::new(), |mut acc, tok| {
2630            acc.extend(tok);
2631            acc
2632        });
2633
2634        assert_eq!(expected.to_string(), generated.to_string());
2635        Ok(())
2636    }
2637
2638    #[test]
2639    fn test_gen_with_attributes() -> io::Result<()> {
2640        let cake_entity = setup().get(0).unwrap().clone();
2641
2642        assert_eq!(cake_entity.get_table_name_snake_case(), "cake");
2643
2644        // Compact code blocks
2645        assert_eq!(
2646            comparable_file_string(include_str!(
2647                "../../tests/compact_with_attributes/cake_none.rs"
2648            ))?,
2649            generated_to_string(EntityWriter::gen_compact_code_blocks(
2650                &cake_entity,
2651                &WithSerde::None,
2652                &DateTimeCrate::Chrono,
2653                &None,
2654                false,
2655                false,
2656                &TokenStream::new(),
2657                &TokenStream::new(),
2658                false,
2659                true,
2660            ))
2661        );
2662        assert_eq!(
2663            comparable_file_string(include_str!(
2664                "../../tests/compact_with_attributes/cake_one.rs"
2665            ))?,
2666            generated_to_string(EntityWriter::gen_compact_code_blocks(
2667                &cake_entity,
2668                &WithSerde::None,
2669                &DateTimeCrate::Chrono,
2670                &None,
2671                false,
2672                false,
2673                &TokenStream::new(),
2674                &bonus_attributes([r#"serde(rename_all = "camelCase")"#]),
2675                false,
2676                true,
2677            ))
2678        );
2679        assert_eq!(
2680            comparable_file_string(include_str!(
2681                "../../tests/compact_with_attributes/cake_multiple.rs"
2682            ))?,
2683            generated_to_string(EntityWriter::gen_compact_code_blocks(
2684                &cake_entity,
2685                &WithSerde::None,
2686                &DateTimeCrate::Chrono,
2687                &None,
2688                false,
2689                false,
2690                &TokenStream::new(),
2691                &bonus_attributes([r#"serde(rename_all = "camelCase")"#, "ts(export)"]),
2692                false,
2693                true,
2694            ))
2695        );
2696
2697        // Expanded code blocks
2698        assert_eq!(
2699            comparable_file_string(include_str!(
2700                "../../tests/expanded_with_attributes/cake_none.rs"
2701            ))?,
2702            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2703                &cake_entity,
2704                &WithSerde::None,
2705                &DateTimeCrate::Chrono,
2706                &None,
2707                false,
2708                false,
2709                &TokenStream::new(),
2710                &TokenStream::new(),
2711                false,
2712                true,
2713            ))
2714        );
2715        assert_eq!(
2716            comparable_file_string(include_str!(
2717                "../../tests/expanded_with_attributes/cake_one.rs"
2718            ))?,
2719            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2720                &cake_entity,
2721                &WithSerde::None,
2722                &DateTimeCrate::Chrono,
2723                &None,
2724                false,
2725                false,
2726                &TokenStream::new(),
2727                &bonus_attributes([r#"serde(rename_all = "camelCase")"#]),
2728                false,
2729                true,
2730            ))
2731        );
2732        assert_eq!(
2733            comparable_file_string(include_str!(
2734                "../../tests/expanded_with_attributes/cake_multiple.rs"
2735            ))?,
2736            generated_to_string(EntityWriter::gen_expanded_code_blocks(
2737                &cake_entity,
2738                &WithSerde::None,
2739                &DateTimeCrate::Chrono,
2740                &None,
2741                false,
2742                false,
2743                &TokenStream::new(),
2744                &bonus_attributes([r#"serde(rename_all = "camelCase")"#, "ts(export)"]),
2745                false,
2746                true,
2747            ))
2748        );
2749
2750        // Frontend code blocks
2751        assert_eq!(
2752            comparable_file_string(include_str!(
2753                "../../tests/frontend_with_attributes/cake_none.rs"
2754            ))?,
2755            generated_to_string(EntityWriter::gen_frontend_code_blocks(
2756                &cake_entity,
2757                &WithSerde::None,
2758                &DateTimeCrate::Chrono,
2759                &None,
2760                false,
2761                false,
2762                &TokenStream::new(),
2763                &TokenStream::new(),
2764                false,
2765                true,
2766            ))
2767        );
2768        assert_eq!(
2769            comparable_file_string(include_str!(
2770                "../../tests/frontend_with_attributes/cake_one.rs"
2771            ))?,
2772            generated_to_string(EntityWriter::gen_frontend_code_blocks(
2773                &cake_entity,
2774                &WithSerde::None,
2775                &DateTimeCrate::Chrono,
2776                &None,
2777                false,
2778                false,
2779                &TokenStream::new(),
2780                &bonus_attributes([r#"serde(rename_all = "camelCase")"#]),
2781                false,
2782                true,
2783            ))
2784        );
2785        assert_eq!(
2786            comparable_file_string(include_str!(
2787                "../../tests/frontend_with_attributes/cake_multiple.rs"
2788            ))?,
2789            generated_to_string(EntityWriter::gen_frontend_code_blocks(
2790                &cake_entity,
2791                &WithSerde::None,
2792                &DateTimeCrate::Chrono,
2793                &None,
2794                false,
2795                false,
2796                &TokenStream::new(),
2797                &bonus_attributes([r#"serde(rename_all = "camelCase")"#, "ts(export)"]),
2798                false,
2799                true,
2800            ))
2801        );
2802
2803        Ok(())
2804    }
2805
2806    fn generated_to_string(generated: Vec<TokenStream>) -> String {
2807        generated
2808            .into_iter()
2809            .fold(TokenStream::new(), |mut acc, tok| {
2810                acc.extend(tok);
2811                acc
2812            })
2813            .to_string()
2814    }
2815
2816    fn comparable_file_string(file: &str) -> io::Result<String> {
2817        let mut reader = BufReader::new(file.as_bytes());
2818        let mut lines: Vec<String> = Vec::new();
2819
2820        reader.read_until(b'\n', &mut Vec::new())?;
2821
2822        let mut line = String::new();
2823        while reader.read_line(&mut line)? > 0 {
2824            lines.push(line.to_owned());
2825            line.clear();
2826        }
2827        let content = lines.join("");
2828        let expected: TokenStream = content.parse().unwrap();
2829
2830        Ok(expected.to_string())
2831    }
2832
2833    #[test]
2834    fn test_gen_postgres() -> io::Result<()> {
2835        let entities = vec![
2836            // This tests that the JsonBinary column type is annotated
2837            // correctly in compact entity form. More information can be found
2838            // in this issue:
2839            //
2840            // https://github.com/SeaQL/sea-orm/issues/1344
2841            Entity {
2842                table_name: "task".to_owned(),
2843                columns: vec![
2844                    Column {
2845                        name: "id".to_owned(),
2846                        col_type: ColumnType::Integer,
2847                        auto_increment: true,
2848                        not_null: true,
2849                        unique: false,
2850                    },
2851                    Column {
2852                        name: "payload".to_owned(),
2853                        col_type: ColumnType::Json,
2854                        auto_increment: false,
2855                        not_null: true,
2856                        unique: false,
2857                    },
2858                    Column {
2859                        name: "payload_binary".to_owned(),
2860                        col_type: ColumnType::JsonBinary,
2861                        auto_increment: false,
2862                        not_null: true,
2863                        unique: false,
2864                    },
2865                ],
2866                relations: vec![],
2867                conjunct_relations: vec![],
2868                primary_keys: vec![PrimaryKey {
2869                    name: "id".to_owned(),
2870                }],
2871            },
2872        ];
2873        const ENTITY_FILES: [&str; 1] = [include_str!("../../tests/postgres/binary_json.rs")];
2874
2875        const ENTITY_FILES_EXPANDED: [&str; 1] =
2876            [include_str!("../../tests/postgres/binary_json_expanded.rs")];
2877
2878        assert_eq!(entities.len(), ENTITY_FILES.len());
2879
2880        for (i, entity) in entities.iter().enumerate() {
2881            assert_eq!(
2882                parse_from_file(ENTITY_FILES[i].as_bytes())?.to_string(),
2883                EntityWriter::gen_compact_code_blocks(
2884                    entity,
2885                    &crate::WithSerde::None,
2886                    &crate::DateTimeCrate::Chrono,
2887                    &None,
2888                    false,
2889                    false,
2890                    &TokenStream::new(),
2891                    &TokenStream::new(),
2892                    false,
2893                    true,
2894                )
2895                .into_iter()
2896                .skip(1)
2897                .fold(TokenStream::new(), |mut acc, tok| {
2898                    acc.extend(tok);
2899                    acc
2900                })
2901                .to_string()
2902            );
2903            assert_eq!(
2904                parse_from_file(ENTITY_FILES_EXPANDED[i].as_bytes())?.to_string(),
2905                EntityWriter::gen_expanded_code_blocks(
2906                    entity,
2907                    &crate::WithSerde::None,
2908                    &crate::DateTimeCrate::Chrono,
2909                    &Some("schema_name".to_owned()),
2910                    false,
2911                    false,
2912                    &TokenStream::new(),
2913                    &TokenStream::new(),
2914                    false,
2915                    true,
2916                )
2917                .into_iter()
2918                .skip(1)
2919                .fold(TokenStream::new(), |mut acc, tok| {
2920                    acc.extend(tok);
2921                    acc
2922                })
2923                .to_string()
2924            );
2925        }
2926
2927        Ok(())
2928    }
2929
2930    #[test]
2931    fn test_gen_import_active_enum() -> io::Result<()> {
2932        let entities = vec![
2933            Entity {
2934                table_name: "tea_pairing".to_owned(),
2935                columns: vec![
2936                    Column {
2937                        name: "id".to_owned(),
2938                        col_type: ColumnType::Integer,
2939                        auto_increment: true,
2940                        not_null: true,
2941                        unique: false,
2942                    },
2943                    Column {
2944                        name: "first_tea".to_owned(),
2945                        col_type: ColumnType::Enum {
2946                            name: SeaRc::new(Alias::new("tea_enum")),
2947                            variants: vec![
2948                                SeaRc::new(Alias::new("everyday_tea")),
2949                                SeaRc::new(Alias::new("breakfast_tea")),
2950                            ],
2951                        },
2952                        auto_increment: false,
2953                        not_null: true,
2954                        unique: false,
2955                    },
2956                    Column {
2957                        name: "second_tea".to_owned(),
2958                        col_type: ColumnType::Enum {
2959                            name: SeaRc::new(Alias::new("tea_enum")),
2960                            variants: vec![
2961                                SeaRc::new(Alias::new("everyday_tea")),
2962                                SeaRc::new(Alias::new("breakfast_tea")),
2963                            ],
2964                        },
2965                        auto_increment: false,
2966                        not_null: true,
2967                        unique: false,
2968                    },
2969                ],
2970                relations: vec![],
2971                conjunct_relations: vec![],
2972                primary_keys: vec![PrimaryKey {
2973                    name: "id".to_owned(),
2974                }],
2975            },
2976            Entity {
2977                table_name: "tea_pairing_with_size".to_owned(),
2978                columns: vec![
2979                    Column {
2980                        name: "id".to_owned(),
2981                        col_type: ColumnType::Integer,
2982                        auto_increment: true,
2983                        not_null: true,
2984                        unique: false,
2985                    },
2986                    Column {
2987                        name: "first_tea".to_owned(),
2988                        col_type: ColumnType::Enum {
2989                            name: SeaRc::new(Alias::new("tea_enum")),
2990                            variants: vec![
2991                                SeaRc::new(Alias::new("everyday_tea")),
2992                                SeaRc::new(Alias::new("breakfast_tea")),
2993                            ],
2994                        },
2995                        auto_increment: false,
2996                        not_null: true,
2997                        unique: false,
2998                    },
2999                    Column {
3000                        name: "second_tea".to_owned(),
3001                        col_type: ColumnType::Enum {
3002                            name: SeaRc::new(Alias::new("tea_enum")),
3003                            variants: vec![
3004                                SeaRc::new(Alias::new("everyday_tea")),
3005                                SeaRc::new(Alias::new("breakfast_tea")),
3006                            ],
3007                        },
3008                        auto_increment: false,
3009                        not_null: true,
3010                        unique: false,
3011                    },
3012                    Column {
3013                        name: "size".to_owned(),
3014                        col_type: ColumnType::Enum {
3015                            name: SeaRc::new(Alias::new("tea_size")),
3016                            variants: vec![
3017                                SeaRc::new(Alias::new("small")),
3018                                SeaRc::new(Alias::new("medium")),
3019                                SeaRc::new(Alias::new("huge")),
3020                            ],
3021                        },
3022                        auto_increment: false,
3023                        not_null: true,
3024                        unique: false,
3025                    },
3026                ],
3027                relations: vec![],
3028                conjunct_relations: vec![],
3029                primary_keys: vec![PrimaryKey {
3030                    name: "id".to_owned(),
3031                }],
3032            },
3033        ];
3034
3035        assert_eq!(
3036            quote!(
3037                use super::sea_orm_active_enums::TeaEnum;
3038            )
3039            .to_string(),
3040            EntityWriter::gen_import_active_enum(&entities[0]).to_string()
3041        );
3042
3043        assert_eq!(
3044            quote!(
3045                use super::sea_orm_active_enums::TeaEnum;
3046                use super::sea_orm_active_enums::TeaSize;
3047            )
3048            .to_string(),
3049            EntityWriter::gen_import_active_enum(&entities[1]).to_string()
3050        );
3051
3052        Ok(())
3053    }
3054}