From 805f80c36a0a518ae1d918bf4a797a094673cd83 Mon Sep 17 00:00:00 2001
From: YunaiV <zhijiantianya@gmail.com>
Date: Thu, 19 Sep 2024 13:10:23 +0800
Subject: [PATCH] =?UTF-8?q?=E3=80=90=E5=8A=9F=E8=83=BD=E4=BC=98=E5=8C=96?=
 =?UTF-8?q?=E3=80=91INFRA=EF=BC=9A=E4=BB=A3=E7=A0=81=E7=94=9F=E6=88=90?=
 =?UTF-8?q?=E6=97=B6=EF=BC=8C=E9=BB=98=E8=AE=A4=E7=A6=81=E7=94=A8=E5=8D=95?=
 =?UTF-8?q?=E5=85=83=E6=B5=8B=E8=AF=95=E7=9A=84=E7=94=9F=E6=88=90=EF=BC=8C?=
 =?UTF-8?q?=E6=9B=B4=E7=AC=A6=E5=90=88=E5=A4=A7=E5=AE=B6=E7=9A=84=E4=B9=A0?=
 =?UTF-8?q?=E6=83=AF?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../infra/framework/codegen/config/CodegenProperties.java   | 6 ++++++
 .../module/infra/service/codegen/inner/CodegenEngine.java   | 5 +++++
 yudao-server/src/main/resources/application.yaml            | 1 +
 3 files changed, 12 insertions(+)

diff --git a/yudao-module-infra/yudao-module-infra-biz/src/main/java/cn/iocoder/yudao/module/infra/framework/codegen/config/CodegenProperties.java b/yudao-module-infra/yudao-module-infra-biz/src/main/java/cn/iocoder/yudao/module/infra/framework/codegen/config/CodegenProperties.java
index 7cc849f96..cfa1fa2fb 100644
--- a/yudao-module-infra/yudao-module-infra-biz/src/main/java/cn/iocoder/yudao/module/infra/framework/codegen/config/CodegenProperties.java
+++ b/yudao-module-infra/yudao-module-infra-biz/src/main/java/cn/iocoder/yudao/module/infra/framework/codegen/config/CodegenProperties.java
@@ -34,4 +34,10 @@ public class CodegenProperties {
     @NotNull(message = "代码生成的前端类型不能为空")
     private Integer frontType;
 
+    /**
+     * 是否生成单元测试
+     */
+    @NotNull(message = "是否生成单元测试不能为空")
+    private Boolean unitTestEnable;
+
 }
diff --git a/yudao-module-infra/yudao-module-infra-biz/src/main/java/cn/iocoder/yudao/module/infra/service/codegen/inner/CodegenEngine.java b/yudao-module-infra/yudao-module-infra-biz/src/main/java/cn/iocoder/yudao/module/infra/service/codegen/inner/CodegenEngine.java
index 63e0c92ac..8f0068273 100644
--- a/yudao-module-infra/yudao-module-infra-biz/src/main/java/cn/iocoder/yudao/module/infra/service/codegen/inner/CodegenEngine.java
+++ b/yudao-module-infra/yudao-module-infra-biz/src/main/java/cn/iocoder/yudao/module/infra/service/codegen/inner/CodegenEngine.java
@@ -398,6 +398,11 @@ public class CodegenEngine {
         Map<String, String> templates = new LinkedHashMap<>();
         templates.putAll(SERVER_TEMPLATES);
         templates.putAll(FRONT_TEMPLATES.row(frontType));
+        // 如果禁用单元测试，则移除对应的模版
+        if (Boolean.FALSE.equals(codegenProperties.getUnitTestEnable())) {
+            templates.remove(javaTemplatePath("test/serviceTest"));
+            templates.remove("codegen/sql/h2.vm");
+        }
         return templates;
     }
 
diff --git a/yudao-server/src/main/resources/application.yaml b/yudao-server/src/main/resources/application.yaml
index baf68657e..e0e135c7b 100644
--- a/yudao-server/src/main/resources/application.yaml
+++ b/yudao-server/src/main/resources/application.yaml
@@ -246,6 +246,7 @@ yudao:
     base-package: ${yudao.info.base-package}
     db-schemas: ${spring.datasource.dynamic.datasource.master.name}
     front-type: 20 # 前端模版的类型，参见 CodegenFrontTypeEnum 枚举类
+    unit-test-enable: false # 是否生成单元测试
   tenant: # 多租户相关配置项
     enable: true
     ignore-urls: