diff --git a/mybatis-flex-codegen/src/main/java/com/mybatisflex/codegen/entity/Table.java b/mybatis-flex-codegen/src/main/java/com/mybatisflex/codegen/entity/Table.java
index 520a10e0..675c85eb 100644
--- a/mybatis-flex-codegen/src/main/java/com/mybatisflex/codegen/entity/Table.java
+++ b/mybatis-flex-codegen/src/main/java/com/mybatisflex/codegen/entity/Table.java
@@ -128,18 +128,6 @@ public class Table {
 
         EntityConfig entityConfig = globalConfig.getEntityConfig();
 
-        //开启 lombok
-        if (entityConfig.isWithLombok()) {
-            //import lombok.AllArgsConstructor;
-            //import lombok.Builder;
-            //import lombok.Data;
-            //import lombok.NoArgsConstructor;
-            imports.add("lombok.AllArgsConstructor");
-            imports.add("lombok.Builder");
-            imports.add("lombok.Data");
-            imports.add("lombok.NoArgsConstructor");
-        }
-
         if (entityConfig.getSupperClass() != null) {
             imports.add(entityConfig.getSupperClass().getName());
         }
@@ -277,16 +265,6 @@ public class Table {
      */
     public String buildTableAnnotation() {
         StringBuilder tableAnnotation = new StringBuilder();
-        if (globalConfig.getEntityConfig().isWithLombok()) {
-            //@Data
-            //@Builder
-            //@NoArgsConstructor
-            //@AllArgsConstructor
-            tableAnnotation.append("@Data\n");
-            tableAnnotation.append("@Builder\n");
-            tableAnnotation.append("@NoArgsConstructor\n");
-            tableAnnotation.append("@AllArgsConstructor\n");
-        }
 
         tableAnnotation.append("@Table(value = \"").append(name).append("\"");
 
diff --git a/mybatis-flex-codegen/src/main/resources/templates/enjoy/entity.tpl b/mybatis-flex-codegen/src/main/resources/templates/enjoy/entity.tpl
index 14fbd65b..02b40fe3 100644
--- a/mybatis-flex-codegen/src/main/resources/templates/enjoy/entity.tpl
+++ b/mybatis-flex-codegen/src/main/resources/templates/enjoy/entity.tpl
@@ -1,8 +1,15 @@
+#set(withLombok = entityConfig.isWithLombok())
 package #(packageConfig.entityPackage);
 
-#for(importClass:table.buildImports())
+#for(importClass : table.buildImports())
 import #(importClass);
 #end
+#if(withLombok)
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+import lombok.NoArgsConstructor;
+#end
 
 /**
  * #(table.getComment()) 实体类。
@@ -10,15 +17,21 @@ import #(importClass);
  * @author #(javadocConfig.getAuthor())
  * @since #(javadocConfig.getSince())
  */
+#if(withLombok)
+@Data
+@Builder
+@NoArgsConstructor
+@AllArgsConstructor
+#end
 #(table.buildTableAnnotation())
 public class #(table.buildEntityClassName())#(table.buildExtends())#(table.buildImplements()) {
-#for(column: table.columns)
+#for(column : table.columns)
 
     #(column.buildComment())
     #(column.buildAnnotations())private #(column.propertySimpleType) #(column.property);
 #end
 
-    #if(!entityConfig.isWithLombok())
+#if(!withLombok)
     #for(column: table.columns)
     public #(column.propertySimpleType) #(column.getterMethod())() {
         return #(column.property);
@@ -29,5 +42,5 @@ public class #(table.buildEntityClassName())#(table.buildExtends())#(table.build
     }
 
     #end
-    #end
+#end
 }
diff --git a/mybatis-flex-codegen/src/test/java/com/mybatisflex/codegen/test/GeneratorTest.java b/mybatis-flex-codegen/src/test/java/com/mybatisflex/codegen/test/GeneratorTest.java
index bf56f2d8..0c427375 100644
--- a/mybatis-flex-codegen/src/test/java/com/mybatisflex/codegen/test/GeneratorTest.java
+++ b/mybatis-flex-codegen/src/test/java/com/mybatisflex/codegen/test/GeneratorTest.java
@@ -20,6 +20,7 @@ import com.mybatisflex.codegen.Generator;
 import com.mybatisflex.codegen.config.GlobalConfig;
 import com.mybatisflex.spring.service.impl.CacheableServiceImpl;
 import com.zaxxer.hikari.HikariDataSource;
+import org.junit.Test;
 
 import java.util.function.UnaryOperator;
 
@@ -82,7 +83,7 @@ public class GeneratorTest {
         generator.generate();
     }
 
-//    @Test
+    @Test
     public void testCodeGen2() {
         //配置数据源
         HikariDataSource dataSource = new HikariDataSource();
@@ -112,8 +113,8 @@ public class GeneratorTest {
                 .setGenerateTable("sys_user");
 
         //设置模板路径
-        globalConfig.getTemplateConfig()
-                .setEntity("D:\\Documents\\配置文件\\entity.tpl");
+        //globalConfig.getTemplateConfig()
+        //        .setEntity("D:\\Documents\\配置文件\\entity.tpl");
 
         //配置生成 entity
         globalConfig.enableEntity()