后端重构

2026-01-14 14:34:49 +08:00
parent f5f77b32fe
commit 30842b42e5
274 changed files with 39754 additions and 1005 deletions
@@ -0,0 +1,618 @@
+package autocode
+
+import (
+	"fmt"
+	"go/ast"
+	"go/format"
+	"go/parser"
+	"go/token"
+	"io"
+	"os"
+
+	"github.com/pkg/errors"
+)
+
+// WireProviderAst Wire Provider AST 注入结构体
+// 用于向各层的 ProviderSet 中注入新的 Provider
+type WireProviderAst struct {
+	Layer        string // 层名称: biz, data, handler, service, router
+	PackageName  string // 包名
+	ProviderName string // Provider 名称，如 NewArticleUsecase
+	Path         string // 目标文件路径
+}
+
+// NewWireProviderAst 创建 Wire Provider AST 注入器
+func NewWireProviderAst(layer, packageName, providerName, path string) *WireProviderAst {
+	return &WireProviderAst{
+		Layer:        layer,
+		PackageName:  packageName,
+		ProviderName: providerName,
+		Path:         path,
+	}
+}
+
+// Parse 解析文件
+func (a *WireProviderAst) Parse(filename string, writer io.Writer) (*ast.File, error) {
+	if filename == "" {
+		filename = a.Path
+	}
+	if filename == "" {
+		return nil, fmt.Errorf("文件路径为空")
+	}
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
+	if err != nil {
+		return nil, fmt.Errorf("解析文件失败 %s: %w", filename, err)
+	}
+	return file, nil
+}
+
+// Rollback 回滚注入
+func (a *WireProviderAst) Rollback(file *ast.File) error {
+	return removeWireProvider(file, a.ProviderName)
+}
+
+// Injection 注入代码
+func (a *WireProviderAst) Injection(file *ast.File) error {
+	return addWireProvider(file, a.ProviderName)
+}
+
+// Format 格式化输出
+func (a *WireProviderAst) Format(filename string, writer io.Writer, file *ast.File) error {
+	if filename == "" {
+		filename = a.Path
+	}
+	fset := token.NewFileSet()
+	if writer == nil {
+		f, err := os.OpenFile(filename, os.O_WRONLY|os.O_TRUNC, 0666)
+		if err != nil {
+			return errors.Wrapf(err, "打开文件失败: %s", filename)
+		}
+		defer f.Close()
+		writer = f
+	}
+	return format.Node(writer, fset, file)
+}
+
+// addWireProvider 向 wire.NewSet 中添加 Provider
+func addWireProvider(file *ast.File, providerName string) error {
+	var found bool
+	ast.Inspect(file, func(node ast.Node) bool {
+		// 查找 var ProviderSet = wire.NewSet(...) 声明
+		genDecl, ok := node.(*ast.GenDecl)
+		if !ok || genDecl.Tok != token.VAR {
+			return true
+		}
+
+		for _, spec := range genDecl.Specs {
+			valueSpec, ok := spec.(*ast.ValueSpec)
+			if !ok {
+				continue
+			}
+
+			// 检查是否是 ProviderSet 变量
+			for i, name := range valueSpec.Names {
+				if name.Name != "ProviderSet" {
+					continue
+				}
+
+				if i >= len(valueSpec.Values) {
+					continue
+				}
+
+				// 查找 wire.NewSet 调用
+				callExpr, ok := valueSpec.Values[i].(*ast.CallExpr)
+				if !ok {
+					continue
+				}
+
+				selExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
+				if !ok {
+					continue
+				}
+
+				ident, ok := selExpr.X.(*ast.Ident)
+				if !ok || ident.Name != "wire" || selExpr.Sel.Name != "NewSet" {
+					continue
+				}
+
+				// 检查 Provider 是否已存在
+				for _, arg := range callExpr.Args {
+					if argIdent, ok := arg.(*ast.Ident); ok {
+						if argIdent.Name == providerName {
+							found = true
+							return false // 已存在，不需要添加
+						}
+					}
+				}
+
+				// 添加新的 Provider
+				newArg := &ast.Ident{Name: providerName}
+				callExpr.Args = append(callExpr.Args, newArg)
+				found = true
+				return false
+			}
+		}
+		return true
+	})
+
+	if !found {
+		return fmt.Errorf("未找到 ProviderSet 变量或 wire.NewSet 调用")
+	}
+	return nil
+}
+
+// removeWireProvider 从 wire.NewSet 中移除 Provider
+func removeWireProvider(file *ast.File, providerName string) error {
+	ast.Inspect(file, func(node ast.Node) bool {
+		genDecl, ok := node.(*ast.GenDecl)
+		if !ok || genDecl.Tok != token.VAR {
+			return true
+		}
+
+		for _, spec := range genDecl.Specs {
+			valueSpec, ok := spec.(*ast.ValueSpec)
+			if !ok {
+				continue
+			}
+
+			for i, name := range valueSpec.Names {
+				if name.Name != "ProviderSet" {
+					continue
+				}
+
+				if i >= len(valueSpec.Values) {
+					continue
+				}
+
+				callExpr, ok := valueSpec.Values[i].(*ast.CallExpr)
+				if !ok {
+					continue
+				}
+
+				selExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
+				if !ok {
+					continue
+				}
+
+				ident, ok := selExpr.X.(*ast.Ident)
+				if !ok || ident.Name != "wire" || selExpr.Sel.Name != "NewSet" {
+					continue
+				}
+
+				// 移除指定的 Provider
+				newArgs := make([]ast.Expr, 0, len(callExpr.Args))
+				for _, arg := range callExpr.Args {
+					if argIdent, ok := arg.(*ast.Ident); ok {
+						if argIdent.Name == providerName {
+							continue // 跳过要移除的 Provider
+						}
+					}
+					newArgs = append(newArgs, arg)
+				}
+				callExpr.Args = newArgs
+				return false
+			}
+		}
+		return true
+	})
+	return nil
+}
+
+// MainProviderSetAst 主 ProviderSet AST 注入结构体
+// 用于向主层的 ProviderSet（如 internal/biz/biz.go）中注入新包的 Provider
+type MainProviderSetAst struct {
+	Layer        string // 层名称: biz, data, service
+	PackageName  string // 新包名
+	ProviderName string // Provider 名称，如 NewArticleUsecase
+	ImportPath   string // 导入路径
+	Path         string // 目标文件路径
+}
+
+// NewMainProviderSetAst 创建主 ProviderSet AST 注入器
+func NewMainProviderSetAst(layer, packageName, providerName, importPath, path string) *MainProviderSetAst {
+	return &MainProviderSetAst{
+		Layer:        layer,
+		PackageName:  packageName,
+		ProviderName: providerName,
+		ImportPath:   importPath,
+		Path:         path,
+	}
+}
+
+// Parse 解析文件
+func (a *MainProviderSetAst) Parse(filename string, writer io.Writer) (*ast.File, error) {
+	if filename == "" {
+		filename = a.Path
+	}
+	if filename == "" {
+		return nil, fmt.Errorf("文件路径为空")
+	}
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
+	if err != nil {
+		return nil, fmt.Errorf("解析文件失败 %s: %w", filename, err)
+	}
+	return file, nil
+}
+
+// Rollback 回滚注入
+func (a *MainProviderSetAst) Rollback(file *ast.File) error {
+	// 移除导入
+	removeImport(file, a.ImportPath)
+	// 移除 Provider
+	return removeWireProviderWithPackage(file, a.PackageName, a.ProviderName)
+}
+
+// Injection 注入代码
+func (a *MainProviderSetAst) Injection(file *ast.File) error {
+	// 添加导入
+	addImport(file, a.ImportPath)
+	// 添加 Provider
+	return addWireProviderWithPackage(file, a.PackageName, a.ProviderName)
+}
+
+// Format 格式化输出
+func (a *MainProviderSetAst) Format(filename string, writer io.Writer, file *ast.File) error {
+	if filename == "" {
+		filename = a.Path
+	}
+	fset := token.NewFileSet()
+	if writer == nil {
+		f, err := os.OpenFile(filename, os.O_WRONLY|os.O_TRUNC, 0666)
+		if err != nil {
+			return errors.Wrapf(err, "打开文件失败: %s", filename)
+		}
+		defer f.Close()
+		writer = f
+	}
+	return format.Node(writer, fset, file)
+}
+
+// addImport 添加导入
+func addImport(file *ast.File, importPath string) {
+	if importPath == "" {
+		return
+	}
+	impStr := fmt.Sprintf("\"%s\"", importPath)
+
+	// 检查是否已存在
+	for _, imp := range file.Imports {
+		if imp.Path.Value == impStr {
+			return
+		}
+	}
+
+	// 查找 import 声明并添加
+	for _, decl := range file.Decls {
+		genDecl, ok := decl.(*ast.GenDecl)
+		if !ok || genDecl.Tok != token.IMPORT {
+			continue
+		}
+		genDecl.Specs = append(genDecl.Specs, &ast.ImportSpec{
+			Path: &ast.BasicLit{
+				Kind:  token.STRING,
+				Value: impStr,
+			},
+		})
+		return
+	}
+
+	// 如果没有 import 声明，创建一个新的
+	newDecl := &ast.GenDecl{
+		Tok: token.IMPORT,
+		Specs: []ast.Spec{
+			&ast.ImportSpec{
+				Path: &ast.BasicLit{
+					Kind:  token.STRING,
+					Value: impStr,
+				},
+			},
+		},
+	}
+	file.Decls = append([]ast.Decl{newDecl}, file.Decls...)
+}
+
+// removeImport 移除导入
+func removeImport(file *ast.File, importPath string) {
+	if importPath == "" {
+		return
+	}
+	impStr := fmt.Sprintf("\"%s\"", importPath)
+
+	for i := 0; i < len(file.Decls); i++ {
+		genDecl, ok := file.Decls[i].(*ast.GenDecl)
+		if !ok || genDecl.Tok != token.IMPORT {
+			continue
+		}
+		for j := 0; j < len(genDecl.Specs); j++ {
+			impSpec, ok := genDecl.Specs[j].(*ast.ImportSpec)
+			if !ok {
+				continue
+			}
+			if impSpec.Path.Value == impStr {
+				genDecl.Specs = append(genDecl.Specs[:j], genDecl.Specs[j+1:]...)
+				if len(genDecl.Specs) == 0 {
+					file.Decls = append(file.Decls[:i], file.Decls[i+1:]...)
+				}
+				return
+			}
+		}
+	}
+}
+
+// addWireProviderWithPackage 向 wire.NewSet 中添加带包名的 Provider
+// 例如: article.NewArticleUsecase
+func addWireProviderWithPackage(file *ast.File, packageName, providerName string) error {
+	var found bool
+	ast.Inspect(file, func(node ast.Node) bool {
+		genDecl, ok := node.(*ast.GenDecl)
+		if !ok || genDecl.Tok != token.VAR {
+			return true
+		}
+
+		for _, spec := range genDecl.Specs {
+			valueSpec, ok := spec.(*ast.ValueSpec)
+			if !ok {
+				continue
+			}
+
+			for i, name := range valueSpec.Names {
+				if name.Name != "ProviderSet" {
+					continue
+				}
+
+				if i >= len(valueSpec.Values) {
+					continue
+				}
+
+				callExpr, ok := valueSpec.Values[i].(*ast.CallExpr)
+				if !ok {
+					continue
+				}
+
+				selExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
+				if !ok {
+					continue
+				}
+
+				ident, ok := selExpr.X.(*ast.Ident)
+				if !ok || ident.Name != "wire" || selExpr.Sel.Name != "NewSet" {
+					continue
+				}
+
+				// 构建完整的 Provider 名称
+				fullProviderName := fmt.Sprintf("%s.%s", packageName, providerName)
+
+				// 检查 Provider 是否已存在
+				for _, arg := range callExpr.Args {
+					if selArg, ok := arg.(*ast.SelectorExpr); ok {
+						if pkgIdent, ok := selArg.X.(*ast.Ident); ok {
+							existingName := fmt.Sprintf("%s.%s", pkgIdent.Name, selArg.Sel.Name)
+							if existingName == fullProviderName {
+								found = true
+								return false // 已存在，不需要添加
+							}
+						}
+					}
+				}
+
+				// 添加新的 Provider (package.ProviderName 格式)
+				newArg := &ast.SelectorExpr{
+					X:   &ast.Ident{Name: packageName},
+					Sel: &ast.Ident{Name: providerName},
+				}
+				callExpr.Args = append(callExpr.Args, newArg)
+				found = true
+				return false
+			}
+		}
+		return true
+	})
+
+	if !found {
+		return fmt.Errorf("未找到 ProviderSet 变量或 wire.NewSet 调用")
+	}
+	return nil
+}
+
+// removeWireProviderWithPackage 从 wire.NewSet 中移除带包名的 Provider
+func removeWireProviderWithPackage(file *ast.File, packageName, providerName string) error {
+	fullProviderName := fmt.Sprintf("%s.%s", packageName, providerName)
+
+	ast.Inspect(file, func(node ast.Node) bool {
+		genDecl, ok := node.(*ast.GenDecl)
+		if !ok || genDecl.Tok != token.VAR {
+			return true
+		}
+
+		for _, spec := range genDecl.Specs {
+			valueSpec, ok := spec.(*ast.ValueSpec)
+			if !ok {
+				continue
+			}
+
+			for i, name := range valueSpec.Names {
+				if name.Name != "ProviderSet" {
+					continue
+				}
+
+				if i >= len(valueSpec.Values) {
+					continue
+				}
+
+				callExpr, ok := valueSpec.Values[i].(*ast.CallExpr)
+				if !ok {
+					continue
+				}
+
+				selExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
+				if !ok {
+					continue
+				}
+
+				ident, ok := selExpr.X.(*ast.Ident)
+				if !ok || ident.Name != "wire" || selExpr.Sel.Name != "NewSet" {
+					continue
+				}
+
+				// 移除指定的 Provider
+				newArgs := make([]ast.Expr, 0, len(callExpr.Args))
+				for _, arg := range callExpr.Args {
+					if selArg, ok := arg.(*ast.SelectorExpr); ok {
+						if pkgIdent, ok := selArg.X.(*ast.Ident); ok {
+							existingName := fmt.Sprintf("%s.%s", pkgIdent.Name, selArg.Sel.Name)
+							if existingName == fullProviderName {
+								continue // 跳过要移除的 Provider
+							}
+						}
+					}
+					newArgs = append(newArgs, arg)
+				}
+				callExpr.Args = newArgs
+				return false
+			}
+		}
+		return true
+	})
+	return nil
+}
+
+// GetLayerMainFile 获取层的主文件路径
+func GetLayerMainFile(layer string) string {
+	switch layer {
+	case "biz":
+		return "internal/biz/biz.go"
+	case "data":
+		return "internal/data/data.go"
+	case "service":
+		return "internal/service/service.go"
+	case "handler", "api":
+		return "internal/server/handler/enter.go"
+	case "router":
+		return "internal/server/router/enter.go"
+	default:
+		return ""
+	}
+}
+
+// GetLayerImportPath 获取层的导入路径
+func GetLayerImportPath(module, layer, packageName string) string {
+	switch layer {
+	case "biz":
+		return fmt.Sprintf("%s/internal/biz/%s", module, packageName)
+	case "data":
+		return fmt.Sprintf("%s/internal/data/%s", module, packageName)
+	case "service":
+		return fmt.Sprintf("%s/internal/service/%s", module, packageName)
+	case "handler", "api":
+		return fmt.Sprintf("%s/internal/server/handler/%s", module, packageName)
+	case "router":
+		return fmt.Sprintf("%s/internal/server/router/%s", module, packageName)
+	default:
+		return ""
+	}
+}
+
+// GetProviderName 获取 Provider 名称
+func GetProviderName(layer, structName string) string {
+	switch layer {
+	case "biz":
+		return fmt.Sprintf("New%sUsecase", structName)
+	case "data":
+		return fmt.Sprintf("New%sRepo", structName)
+	case "service":
+		return fmt.Sprintf("New%sService", structName)
+	case "handler", "api":
+		return fmt.Sprintf("New%sHandler", structName)
+	case "router":
+		return fmt.Sprintf("New%sRouter", structName)
+	default:
+		return ""
+	}
+}
+
+// InjectWireProvider 注入 Wire Provider 到主文件
+// 这是一个便捷函数，用于一次性完成导入和 Provider 注入
+func InjectWireProvider(module, layer, packageName, structName string) error {
+	mainFile := GetLayerMainFile(layer)
+	if mainFile == "" {
+		return fmt.Errorf("未知的层: %s", layer)
+	}
+
+	importPath := GetLayerImportPath(module, layer, packageName)
+	providerName := GetProviderName(layer, structName)
+
+	injector := NewMainProviderSetAst(layer, packageName, providerName, importPath, mainFile)
+
+	file, err := injector.Parse("", nil)
+	if err != nil {
+		return err
+	}
+
+	if err := injector.Injection(file); err != nil {
+		return err
+	}
+
+	return injector.Format("", nil, file)
+}
+
+// RollbackWireProvider 回滚 Wire Provider 注入
+func RollbackWireProvider(module, layer, packageName, structName string) error {
+	mainFile := GetLayerMainFile(layer)
+	if mainFile == "" {
+		return fmt.Errorf("未知的层: %s", layer)
+	}
+
+	importPath := GetLayerImportPath(module, layer, packageName)
+	providerName := GetProviderName(layer, structName)
+
+	injector := NewMainProviderSetAst(layer, packageName, providerName, importPath, mainFile)
+
+	file, err := injector.Parse("", nil)
+	if err != nil {
+		return err
+	}
+
+	if err := injector.Rollback(file); err != nil {
+		return err
+	}
+
+	return injector.Format("", nil, file)
+}
+
+// WireInjectionConfig Wire 注入配置
+type WireInjectionConfig struct {
+	Module      string   // 模块名 (go.mod 中的 module)
+	PackageName string   // 包名
+	StructName  string   // 结构体名
+	Layers      []string // 需要注入的层
+}
+
+// InjectAllWireProviders 注入所有层的 Wire Provider
+func InjectAllWireProviders(config *WireInjectionConfig) error {
+	for _, layer := range config.Layers {
+		if err := InjectWireProvider(config.Module, layer, config.PackageName, config.StructName); err != nil {
+			// 如果注入失败，尝试回滚已注入的
+			for _, rollbackLayer := range config.Layers {
+				if rollbackLayer == layer {
+					break
+				}
+				_ = RollbackWireProvider(config.Module, rollbackLayer, config.PackageName, config.StructName)
+			}
+			return fmt.Errorf("注入 %s 层 Wire Provider 失败: %w", layer, err)
+		}
+	}
+	return nil
+}
+
+// RollbackAllWireProviders 回滚所有层的 Wire Provider
+func RollbackAllWireProviders(config *WireInjectionConfig) error {
+	var lastErr error
+	for _, layer := range config.Layers {
+		if err := RollbackWireProvider(config.Module, layer, config.PackageName, config.StructName); err != nil {
+			lastErr = err
+		}
+	}
+	return lastErr
+}
@@ -0,0 +1,684 @@
+package autocode
+
+import (
+	"bytes"
+	"go/ast"
+	"go/parser"
+	"go/token"
+	"strings"
+	"testing"
+)
+
+// TestWireProviderAst_Parse 测试 WireProviderAst 解析功能
+func TestWireProviderAst_Parse(t *testing.T) {
+	// 创建测试用的 Go 代码
+	testCode := `package biz
+
+import "github.com/google/wire"
+
+var ProviderSet = wire.NewSet(NewUserUsecase)
+`
+	// 创建临时文件
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
+	if err != nil {
+		t.Fatalf("解析测试代码失败: %v", err)
+	}
+
+	if file == nil {
+		t.Fatal("解析结果为空")
+	}
+
+	if file.Name.Name != "biz" {
+		t.Errorf("包名应为 'biz', 实际为 '%s'", file.Name.Name)
+	}
+}
+
+// TestWireProviderAst_Injection 测试 WireProviderAst 注入功能
+func TestWireProviderAst_Injection(t *testing.T) {
+	testCode := `package biz
+
+import "github.com/google/wire"
+
+var ProviderSet = wire.NewSet(NewUserUsecase)
+`
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
+	if err != nil {
+		t.Fatalf("解析测试代码失败: %v", err)
+	}
+
+	// 创建注入器
+	injector := &WireProviderAst{
+		Layer:        "biz",
+		PackageName:  "article",
+		ProviderName: "NewArticleUsecase",
+		Path:         "test.go",
+	}
+
+	// 执行注入
+	err = injector.Injection(file)
+	if err != nil {
+		t.Fatalf("注入失败: %v", err)
+	}
+
+	// 验证注入结果
+	var found bool
+	ast.Inspect(file, func(node ast.Node) bool {
+		if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.VAR {
+			for _, spec := range genDecl.Specs {
+				if valueSpec, ok := spec.(*ast.ValueSpec); ok {
+					for i, name := range valueSpec.Names {
+						if name.Name == "ProviderSet" && i < len(valueSpec.Values) {
+							if callExpr, ok := valueSpec.Values[i].(*ast.CallExpr); ok {
+								for _, arg := range callExpr.Args {
+									if ident, ok := arg.(*ast.Ident); ok && ident.Name == "NewArticleUsecase" {
+										found = true
+										return false
+									}
+								}
+							}
+						}
+					}
+				}
+			}
+		}
+		return true
+	})
+
+	if !found {
+		t.Error("注入后应包含 NewArticleUsecase Provider")
+	}
+}
+
+// TestWireProviderAst_Injection_Duplicate 测试重复注入不会添加重复项
+func TestWireProviderAst_Injection_Duplicate(t *testing.T) {
+	testCode := `package biz
+
+import "github.com/google/wire"
+
+var ProviderSet = wire.NewSet(NewUserUsecase, NewArticleUsecase)
+`
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
+	if err != nil {
+		t.Fatalf("解析测试代码失败: %v", err)
+	}
+
+	injector := &WireProviderAst{
+		Layer:        "biz",
+		PackageName:  "article",
+		ProviderName: "NewArticleUsecase",
+		Path:         "test.go",
+	}
+
+	// 执行注入（应该不会添加重复项）
+	err = injector.Injection(file)
+	if err != nil {
+		t.Fatalf("注入失败: %v", err)
+	}
+
+	// 统计 NewArticleUsecase 出现次数
+	count := 0
+	ast.Inspect(file, func(node ast.Node) bool {
+		if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.VAR {
+			for _, spec := range genDecl.Specs {
+				if valueSpec, ok := spec.(*ast.ValueSpec); ok {
+					for i, name := range valueSpec.Names {
+						if name.Name == "ProviderSet" && i < len(valueSpec.Values) {
+							if callExpr, ok := valueSpec.Values[i].(*ast.CallExpr); ok {
+								for _, arg := range callExpr.Args {
+									if ident, ok := arg.(*ast.Ident); ok && ident.Name == "NewArticleUsecase" {
+										count++
+									}
+								}
+							}
+						}
+					}
+				}
+			}
+		}
+		return true
+	})
+
+	if count != 1 {
+		t.Errorf("NewArticleUsecase 应该只出现 1 次, 实际出现 %d 次", count)
+	}
+}
+
+// TestWireProviderAst_Rollback 测试 WireProviderAst 回滚功能
+func TestWireProviderAst_Rollback(t *testing.T) {
+	testCode := `package biz
+
+import "github.com/google/wire"
+
+var ProviderSet = wire.NewSet(NewUserUsecase, NewArticleUsecase)
+`
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
+	if err != nil {
+		t.Fatalf("解析测试代码失败: %v", err)
+	}
+
+	injector := &WireProviderAst{
+		Layer:        "biz",
+		PackageName:  "article",
+		ProviderName: "NewArticleUsecase",
+		Path:         "test.go",
+	}
+
+	// 执行回滚
+	err = injector.Rollback(file)
+	if err != nil {
+		t.Fatalf("回滚失败: %v", err)
+	}
+
+	// 验证回滚结果
+	var found bool
+	ast.Inspect(file, func(node ast.Node) bool {
+		if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.VAR {
+			for _, spec := range genDecl.Specs {
+				if valueSpec, ok := spec.(*ast.ValueSpec); ok {
+					for i, name := range valueSpec.Names {
+						if name.Name == "ProviderSet" && i < len(valueSpec.Values) {
+							if callExpr, ok := valueSpec.Values[i].(*ast.CallExpr); ok {
+								for _, arg := range callExpr.Args {
+									if ident, ok := arg.(*ast.Ident); ok && ident.Name == "NewArticleUsecase" {
+										found = true
+										return false
+									}
+								}
+							}
+						}
+					}
+				}
+			}
+		}
+		return true
+	})
+
+	if found {
+		t.Error("回滚后不应包含 NewArticleUsecase Provider")
+	}
+}
+
+// TestWireProviderAst_Format 测试 WireProviderAst 格式化输出
+func TestWireProviderAst_Format(t *testing.T) {
+	testCode := `package biz
+
+import "github.com/google/wire"
+
+var ProviderSet = wire.NewSet(NewUserUsecase)
+`
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
+	if err != nil {
+		t.Fatalf("解析测试代码失败: %v", err)
+	}
+
+	injector := &WireProviderAst{
+		Layer:        "biz",
+		PackageName:  "article",
+		ProviderName: "NewArticleUsecase",
+		Path:         "test.go",
+	}
+
+	// 注入后格式化
+	_ = injector.Injection(file)
+
+	var buf bytes.Buffer
+	err = injector.Format("", &buf, file)
+	if err != nil {
+		t.Fatalf("格式化失败: %v", err)
+	}
+
+	result := buf.String()
+	if !strings.Contains(result, "NewArticleUsecase") {
+		t.Error("格式化输出应包含 NewArticleUsecase")
+	}
+	if !strings.Contains(result, "package biz") {
+		t.Error("格式化输出应包含 package biz")
+	}
+}
+
+// TestMainProviderSetAst_Injection 测试 MainProviderSetAst 注入功能
+func TestMainProviderSetAst_Injection(t *testing.T) {
+	testCode := `package biz
+
+import (
+	"github.com/google/wire"
+	"kra/internal/biz/user"
+)
+
+var ProviderSet = wire.NewSet(user.NewUserUsecase)
+`
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
+	if err != nil {
+		t.Fatalf("解析测试代码失败: %v", err)
+	}
+
+	injector := &MainProviderSetAst{
+		Layer:        "biz",
+		PackageName:  "article",
+		ProviderName: "NewArticleUsecase",
+		ImportPath:   "kra/internal/biz/article",
+		Path:         "test.go",
+	}
+
+	// 执行注入
+	err = injector.Injection(file)
+	if err != nil {
+		t.Fatalf("注入失败: %v", err)
+	}
+
+	// 验证导入是否添加 - 检查 import 声明中是否包含新导入
+	var importFound bool
+	ast.Inspect(file, func(node ast.Node) bool {
+		if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
+			for _, spec := range genDecl.Specs {
+				if impSpec, ok := spec.(*ast.ImportSpec); ok {
+					if impSpec.Path.Value == "\"kra/internal/biz/article\"" {
+						importFound = true
+						return false
+					}
+				}
+			}
+		}
+		return true
+	})
+	if !importFound {
+		t.Error("注入后应包含 kra/internal/biz/article 导入")
+	}
+
+	// 验证 Provider 是否添加
+	var providerFound bool
+	ast.Inspect(file, func(node ast.Node) bool {
+		if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.VAR {
+			for _, spec := range genDecl.Specs {
+				if valueSpec, ok := spec.(*ast.ValueSpec); ok {
+					for i, name := range valueSpec.Names {
+						if name.Name == "ProviderSet" && i < len(valueSpec.Values) {
+							if callExpr, ok := valueSpec.Values[i].(*ast.CallExpr); ok {
+								for _, arg := range callExpr.Args {
+									if selExpr, ok := arg.(*ast.SelectorExpr); ok {
+										if pkgIdent, ok := selExpr.X.(*ast.Ident); ok {
+											if pkgIdent.Name == "article" && selExpr.Sel.Name == "NewArticleUsecase" {
+												providerFound = true
+												return false
+											}
+										}
+									}
+								}
+							}
+						}
+					}
+				}
+			}
+		}
+		return true
+	})
+
+	if !providerFound {
+		t.Error("注入后应包含 article.NewArticleUsecase Provider")
+	}
+}
+
+// TestMainProviderSetAst_Rollback 测试 MainProviderSetAst 回滚功能
+func TestMainProviderSetAst_Rollback(t *testing.T) {
+	testCode := `package biz
+
+import (
+	"github.com/google/wire"
+	"kra/internal/biz/user"
+	"kra/internal/biz/article"
+)
+
+var ProviderSet = wire.NewSet(user.NewUserUsecase, article.NewArticleUsecase)
+`
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
+	if err != nil {
+		t.Fatalf("解析测试代码失败: %v", err)
+	}
+
+	injector := &MainProviderSetAst{
+		Layer:        "biz",
+		PackageName:  "article",
+		ProviderName: "NewArticleUsecase",
+		ImportPath:   "kra/internal/biz/article",
+		Path:         "test.go",
+	}
+
+	// 执行回滚
+	err = injector.Rollback(file)
+	if err != nil {
+		t.Fatalf("回滚失败: %v", err)
+	}
+
+	// 验证导入是否移除 - 检查 import 声明
+	var importFound bool
+	ast.Inspect(file, func(node ast.Node) bool {
+		if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
+			for _, spec := range genDecl.Specs {
+				if impSpec, ok := spec.(*ast.ImportSpec); ok {
+					if impSpec.Path.Value == "\"kra/internal/biz/article\"" {
+						importFound = true
+						return false
+					}
+				}
+			}
+		}
+		return true
+	})
+	if importFound {
+		t.Error("回滚后不应包含 kra/internal/biz/article 导入")
+	}
+
+	// 验证 Provider 是否移除
+	var providerFound bool
+	ast.Inspect(file, func(node ast.Node) bool {
+		if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.VAR {
+			for _, spec := range genDecl.Specs {
+				if valueSpec, ok := spec.(*ast.ValueSpec); ok {
+					for i, name := range valueSpec.Names {
+						if name.Name == "ProviderSet" && i < len(valueSpec.Values) {
+							if callExpr, ok := valueSpec.Values[i].(*ast.CallExpr); ok {
+								for _, arg := range callExpr.Args {
+									if selExpr, ok := arg.(*ast.SelectorExpr); ok {
+										if pkgIdent, ok := selExpr.X.(*ast.Ident); ok {
+											if pkgIdent.Name == "article" && selExpr.Sel.Name == "NewArticleUsecase" {
+												providerFound = true
+												return false
+											}
+										}
+									}
+								}
+							}
+						}
+					}
+				}
+			}
+		}
+		return true
+	})
+
+	if providerFound {
+		t.Error("回滚后不应包含 article.NewArticleUsecase Provider")
+	}
+}
+
+// TestGetLayerMainFile 测试获取层主文件路径
+func TestGetLayerMainFile(t *testing.T) {
+	testCases := []struct {
+		layer    string
+		expected string
+	}{
+		{"biz", "internal/biz/biz.go"},
+		{"data", "internal/data/data.go"},
+		{"service", "internal/service/service.go"},
+		{"handler", "internal/server/handler/enter.go"},
+		{"api", "internal/server/handler/enter.go"},
+		{"router", "internal/server/router/enter.go"},
+		{"unknown", ""},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.layer, func(t *testing.T) {
+			result := GetLayerMainFile(tc.layer)
+			if result != tc.expected {
+				t.Errorf("GetLayerMainFile(%s) = %s, 期望 %s", tc.layer, result, tc.expected)
+			}
+		})
+	}
+}
+
+// TestGetLayerImportPath 测试获取层导入路径
+func TestGetLayerImportPath(t *testing.T) {
+	module := "kra"
+	packageName := "article"
+
+	testCases := []struct {
+		layer    string
+		expected string
+	}{
+		{"biz", "kra/internal/biz/article"},
+		{"data", "kra/internal/data/article"},
+		{"service", "kra/internal/service/article"},
+		{"handler", "kra/internal/server/handler/article"},
+		{"api", "kra/internal/server/handler/article"},
+		{"router", "kra/internal/server/router/article"},
+		{"unknown", ""},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.layer, func(t *testing.T) {
+			result := GetLayerImportPath(module, tc.layer, packageName)
+			if result != tc.expected {
+				t.Errorf("GetLayerImportPath(%s, %s, %s) = %s, 期望 %s", module, tc.layer, packageName, result, tc.expected)
+			}
+		})
+	}
+}
+
+// TestGetProviderName 测试获取 Provider 名称
+func TestGetProviderName(t *testing.T) {
+	structName := "Article"
+
+	testCases := []struct {
+		layer    string
+		expected string
+	}{
+		{"biz", "NewArticleUsecase"},
+		{"data", "NewArticleRepo"},
+		{"service", "NewArticleService"},
+		{"handler", "NewArticleHandler"},
+		{"api", "NewArticleHandler"},
+		{"router", "NewArticleRouter"},
+		{"unknown", ""},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.layer, func(t *testing.T) {
+			result := GetProviderName(tc.layer, structName)
+			if result != tc.expected {
+				t.Errorf("GetProviderName(%s, %s) = %s, 期望 %s", tc.layer, structName, result, tc.expected)
+			}
+		})
+	}
+}
+
+// TestAddImport 测试添加导入
+func TestAddImport(t *testing.T) {
+	testCode := `package main
+
+import "fmt"
+
+func main() {}
+`
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
+	if err != nil {
+		t.Fatalf("解析测试代码失败: %v", err)
+	}
+
+	// 添加新导入
+	addImport(file, "kra/internal/biz/article")
+
+	// 验证导入是否添加 - 检查 import 声明
+	var found bool
+	ast.Inspect(file, func(node ast.Node) bool {
+		if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
+			for _, spec := range genDecl.Specs {
+				if impSpec, ok := spec.(*ast.ImportSpec); ok {
+					if impSpec.Path.Value == "\"kra/internal/biz/article\"" {
+						found = true
+						return false
+					}
+				}
+			}
+		}
+		return true
+	})
+	if !found {
+		t.Error("应该添加 kra/internal/biz/article 导入")
+	}
+}
+
+// TestAddImport_Duplicate 测试重复添加导入
+func TestAddImport_Duplicate(t *testing.T) {
+	testCode := `package main
+
+import (
+	"fmt"
+	"kra/internal/biz/article"
+)
+
+func main() {}
+`
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
+	if err != nil {
+		t.Fatalf("解析测试代码失败: %v", err)
+	}
+
+	initialCount := len(file.Imports)
+
+	// 尝试添加已存在的导入
+	addImport(file, "kra/internal/biz/article")
+
+	// 验证导入数量没有增加
+	if len(file.Imports) != initialCount {
+		t.Errorf("重复添加导入后数量应保持 %d, 实际为 %d", initialCount, len(file.Imports))
+	}
+}
+
+// TestRemoveImport 测试移除导入
+func TestRemoveImport(t *testing.T) {
+	testCode := `package main
+
+import (
+	"fmt"
+	"kra/internal/biz/article"
+)
+
+func main() {}
+`
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
+	if err != nil {
+		t.Fatalf("解析测试代码失败: %v", err)
+	}
+
+	// 移除导入
+	removeImport(file, "kra/internal/biz/article")
+
+	// 验证导入是否移除 - 检查 import 声明
+	var found bool
+	ast.Inspect(file, func(node ast.Node) bool {
+		if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
+			for _, spec := range genDecl.Specs {
+				if impSpec, ok := spec.(*ast.ImportSpec); ok {
+					if impSpec.Path.Value == "\"kra/internal/biz/article\"" {
+						found = true
+						return false
+					}
+				}
+			}
+		}
+		return true
+	})
+	if found {
+		t.Error("应该移除 kra/internal/biz/article 导入")
+	}
+}
+
+// TestNewWireProviderAst 测试创建 WireProviderAst
+func TestNewWireProviderAst(t *testing.T) {
+	ast := NewWireProviderAst("biz", "article", "NewArticleUsecase", "internal/biz/article/enter.go")
+
+	if ast.Layer != "biz" {
+		t.Errorf("Layer 应为 'biz', 实际为 '%s'", ast.Layer)
+	}
+	if ast.PackageName != "article" {
+		t.Errorf("PackageName 应为 'article', 实际为 '%s'", ast.PackageName)
+	}
+	if ast.ProviderName != "NewArticleUsecase" {
+		t.Errorf("ProviderName 应为 'NewArticleUsecase', 实际为 '%s'", ast.ProviderName)
+	}
+	if ast.Path != "internal/biz/article/enter.go" {
+		t.Errorf("Path 应为 'internal/biz/article/enter.go', 实际为 '%s'", ast.Path)
+	}
+}
+
+// TestNewMainProviderSetAst 测试创建 MainProviderSetAst
+func TestNewMainProviderSetAst(t *testing.T) {
+	ast := NewMainProviderSetAst("biz", "article", "NewArticleUsecase", "kra/internal/biz/article", "internal/biz/biz.go")
+
+	if ast.Layer != "biz" {
+		t.Errorf("Layer 应为 'biz', 实际为 '%s'", ast.Layer)
+	}
+	if ast.PackageName != "article" {
+		t.Errorf("PackageName 应为 'article', 实际为 '%s'", ast.PackageName)
+	}
+	if ast.ProviderName != "NewArticleUsecase" {
+		t.Errorf("ProviderName 应为 'NewArticleUsecase', 实际为 '%s'", ast.ProviderName)
+	}
+	if ast.ImportPath != "kra/internal/biz/article" {
+		t.Errorf("ImportPath 应为 'kra/internal/biz/article', 实际为 '%s'", ast.ImportPath)
+	}
+	if ast.Path != "internal/biz/biz.go" {
+		t.Errorf("Path 应为 'internal/biz/biz.go', 实际为 '%s'", ast.Path)
+	}
+}
+
+// TestWireProviderAst_Injection_NoProviderSet 测试没有 ProviderSet 时的注入
+func TestWireProviderAst_Injection_NoProviderSet(t *testing.T) {
+	testCode := `package biz
+
+import "github.com/google/wire"
+
+var SomeOtherVar = "test"
+`
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
+	if err != nil {
+		t.Fatalf("解析测试代码失败: %v", err)
+	}
+
+	injector := &WireProviderAst{
+		Layer:        "biz",
+		PackageName:  "article",
+		ProviderName: "NewArticleUsecase",
+		Path:         "test.go",
+	}
+
+	// 执行注入应该返回错误
+	err = injector.Injection(file)
+	if err == nil {
+		t.Error("没有 ProviderSet 时注入应该返回错误")
+	}
+}
+
+// TestWireInjectionConfig 测试 WireInjectionConfig 结构体
+func TestWireInjectionConfig(t *testing.T) {
+	config := &WireInjectionConfig{
+		Module:      "kra",
+		PackageName: "article",
+		StructName:  "Article",
+		Layers:      []string{"biz", "data", "service", "handler", "router"},
+	}
+
+	if config.Module != "kra" {
+		t.Errorf("Module 应为 'kra', 实际为 '%s'", config.Module)
+	}
+	if config.PackageName != "article" {
+		t.Errorf("PackageName 应为 'article', 实际为 '%s'", config.PackageName)
+	}
+	if config.StructName != "Article" {
+		t.Errorf("StructName 应为 'Article', 实际为 '%s'", config.StructName)
+	}
+	if len(config.Layers) != 5 {
+		t.Errorf("Layers 长度应为 5, 实际为 %d", len(config.Layers))
+	}
+}