diff --git a/src/tests_cfg/cake.rs b/src/tests_cfg/cake.rs
index 0eeb0738..74b37eaf 100644
--- a/src/tests_cfg/cake.rs
+++ b/src/tests_cfg/cake.rs
@@ -1,56 +1,19 @@
 use crate as sea_orm;
-use crate::entity::prelude::*;
+use sea_orm::entity::prelude::*;
 
-#[derive(Copy, Clone, Default, Debug, DeriveEntity)]
-pub struct Entity;
-
-impl EntityName for Entity {
-    fn table_name(&self) -> &str {
-        "cake"
-    }
-}
-
-#[derive(Clone, Debug, PartialEq, DeriveModel, DeriveActiveModel)]
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel, DeriveModel, DeriveActiveModel)]
+#[sea_orm(table_name = "cake")]
 pub struct Model {
+    #[sea_orm(primary_key)]
     pub id: i32,
     pub name: String,
 }
 
-#[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
-pub enum Column {
-    Id,
-    Name,
-}
-
-#[derive(Copy, Clone, Debug, EnumIter, DerivePrimaryKey)]
-pub enum PrimaryKey {
-    Id,
-}
-
-impl PrimaryKeyTrait for PrimaryKey {
-    type ValueType = i32;
-
-    fn auto_increment() -> bool {
-        true
-    }
-}
-
 #[derive(Copy, Clone, Debug, EnumIter)]
 pub enum Relation {
     Fruit,
 }
 
-impl ColumnTrait for Column {
-    type EntityName = Entity;
-
-    fn def(&self) -> ColumnDef {
-        match self {
-            Self::Id => ColumnType::Integer.def(),
-            Self::Name => ColumnType::String(None).def(),
-        }
-    }
-}
-
 impl RelationTrait for Relation {
     fn def(&self) -> RelationDef {
         match self {