From d169f9b9d1ae80f8bbcdd9dc6ea29e817f0f352e Mon Sep 17 00:00:00 2001
From: huang <1724659546@qq.com>
Date: Sat, 13 Sep 2025 15:14:08 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E8=8E=B7=E5=8F=96=E8=A1=A8?=
 =?UTF-8?q?=E5=88=97=E8=A1=A8=E6=96=B9=E6=B3=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/core/application.go           | 10 +---------
 internal/infra/models/models.go        | 13 +++++++++++++
 internal/infra/repository/main_test.go |  2 +-
 3 files changed, 15 insertions(+), 10 deletions(-)
 create mode 100644 internal/infra/models/models.go

diff --git a/internal/core/application.go b/internal/core/application.go
index 2335f06..3db0fa8 100644
--- a/internal/core/application.go
+++ b/internal/core/application.go
@@ -124,15 +124,7 @@ func initStorage(cfg config.DatabaseConfig, logger *logs.Logger) (database.Stora
 	}
 
 	// 执行数据库迁移
-	// 这里需要添加所有需要自动迁移的模型
-	var dbModels = []interface{}{
-		&models.User{},
-		&models.Device{},
-		&models.Plan{},
-		&models.SubPlan{},
-		&models.Task{},
-	}
-	if err := storage.Migrate(dbModels...); err != nil {
+	if err := storage.Migrate(models.GetAllModels()...); err != nil {
 		return nil, fmt.Errorf("数据库迁移失败: %w", err)
 	}
 
diff --git a/internal/infra/models/models.go b/internal/infra/models/models.go
new file mode 100644
index 0000000..37cc5c5
--- /dev/null
+++ b/internal/infra/models/models.go
@@ -0,0 +1,13 @@
+package models
+
+// GetAllModels 返回一个包含所有数据库模型实例的切片。
+// 这个函数用于在数据库初始化时自动迁移所有的表结构。
+func GetAllModels() []interface{} {
+	return []interface{}{
+		&User{},
+		&Device{},
+		&Plan{},
+		&SubPlan{},
+		&Task{},
+	}
+}
diff --git a/internal/infra/repository/main_test.go b/internal/infra/repository/main_test.go
index 700e438..81c3907 100644
--- a/internal/infra/repository/main_test.go
+++ b/internal/infra/repository/main_test.go
@@ -17,7 +17,7 @@ func setupTestDB(t *testing.T) *gorm.DB {
 	assert.NoError(t, err, "连接内存数据库时发生错误")
 
 	// 自动迁移所有需要的表结构
-	err = db.AutoMigrate(&models.User{}, &models.Device{}, &models.SubPlan{}, &models.Task{}, &models.Plan{})
+	err = db.AutoMigrate(models.GetAllModels()...)
 	assert.NoError(t, err, "数据库迁移时发生错误")
 
 	return db