package ast import ( "go/ast" "go/format" "go/parser" "go/token" "os" "strings" "errors" ) const ( INIT_TEST_GO_FNAME = "init_test.go" ) type ASTer struct { pkg string pkg_test_go_fname string init_test_fset *token.FileSet init_test_node *ast.File pkg_test_fset *token.FileSet pkg_test_node *ast.File } func NewASTer() (*ASTer, error) { a := &ASTer{} pkg, err := получитьИмяGoПакетаВЭтойДире() if err != nil { return nil, err } a.pkg = pkg err = a.найтиТестовыйФайл() if err != nil { return nil, err } return a, nil } func (a *ASTer) ДобавитьШаг(шаг, f, ps string) error { return a.добавитьШаг(шаг, f, ps) } func (a *ASTer) найтиТестовыйФайл() error { a.сгенеритьИмяФайла() _, err := os.Stat(INIT_TEST_GO_FNAME) if err != nil { return err } _, err = os.Stat(a.pkg_test_go_fname) return err } func (a *ASTer) сгенеритьИмяФайла() { a.pkg_test_go_fname = a.pkg + "_test.go" } func (a *ASTer) добавитьШаг(шаг, f, ps string) error { err := a.добавитьШагвInitФайл(шаг, f) if err != nil { return err } err = a.добавитьФункциювТестовыйФайл(f, ps) if err != nil { return err } return nil } func (a *ASTer) добавитьШагвInitФайл(шаг, f string) error { err := a.спарситьInitФайлиДобавитьШаг(шаг, f) if err != nil { return err } err = a.перезаписатьInitФайл() if err != nil { return err } return nil } func (a *ASTer) добавитьФункциювТестовыйФайл(f, ps string) error { err := a.спарситьФайлиДобавитьФункцию(f, ps) if err != nil { return err } err = a.перезаписатьФайл() if err != nil { return err } return nil } func (a *ASTer) спарситьInitФайлиДобавитьШаг(шаг, func_name string) error { err := a.спарситьInitФайл() if err != nil { return err } f := a.создатьШаг(шаг, func_name) a.добавитьШагвИнициализаторЕслиНету(f) return nil } func (a *ASTer) спарситьФайлиДобавитьФункцию(func_name, ps string) error { err := a.спарситьФайл() if err != nil { return err } f := a.создатьФункцию(func_name, ps) a.добавитьФункцию(f) return nil } func (a *ASTer) добавитьШагвИнициализаторЕслиНету(step *ast.ExprStmt) { for _, d := range a.init_test_node.Decls { f, ok := d.(*ast.FuncDecl) if !ok { continue } if f.Name == nil { continue } if f.Name.Name == "InitializeScenario" { a.добавитьШагвФункциюInitializeScenario(f, step) return } } } func (a *ASTer) добавитьШагвФункциюInitializeScenario(f *ast.FuncDecl, step *ast.ExprStmt) { for i, stmt := range f.Body.List { if являетсяЛиШагомсТакойЖеФункцией(stmt, step) { return } if !являетсяЛиШагом(stmt) { n := append([]ast.Stmt{step}, f.Body.List[i:]...) f.Body.List = append(f.Body.List[0:i], n...) return } } } func (a *ASTer) добавитьФункцию(fun *ast.FuncDecl) { for _, d := range a.pkg_test_node.Decls { f, ok := d.(*ast.FuncDecl) if !ok { continue } if f.Name == nil { continue } if f.Name.Name == fun.Name.Name { return } } a.pkg_test_node.Decls = append(a.pkg_test_node.Decls, fun) } func (a *ASTer) спарситьInitФайл() error { fset := token.NewFileSet() node, err := parser.ParseFile(fset, INIT_TEST_GO_FNAME, nil, parser.AllErrors|parser.ParseComments) a.init_test_fset = fset a.init_test_node = node return err } func (a *ASTer) спарситьФайл() error { fset := token.NewFileSet() node, err := parser.ParseFile(fset, a.pkg_test_go_fname, nil, parser.AllErrors|parser.ParseComments) a.pkg_test_fset = fset a.pkg_test_node = node return err } func (a *ASTer) создатьШаг(шаг, func_name string) *ast.ExprStmt { st := &ast.ExprStmt{ X: &ast.CallExpr{ Fun: &ast.SelectorExpr{ X: ast.NewIdent("ctx"), Sel: ast.NewIdent("Step"), }, Args: []ast.Expr{ &ast.BasicLit{ Kind: token.STRING, Value: шаг, }, ast.NewIdent(func_name), }, }, } return st } func (a *ASTer) создатьФункцию(func_name, ps string) *ast.FuncDecl { f := &ast.FuncDecl{ Name: ast.NewIdent(func_name), Type: &ast.FuncType{ Params: &ast.FieldList{}, }, Body: &ast.BlockStmt{}, } a.добавитьПараметрыФункции(f, ps) return f } func (a *ASTer) добавитьПараметрыФункции(f *ast.FuncDecl, ps string) { params := a.сконвертитьСтрокувПараметры(ps) var names []*ast.Ident for _, p := range params { switch len(p) { case 0: continue case 1: names = append(names, ast.NewIdent(p[0])) case 2: names = append(names, ast.NewIdent(p[0])) field := &ast.Field{ Names: names, Type: ast.NewIdent(p[1]), } f.Type.Params.List = append(f.Type.Params.List, field) names = nil } } } func (a *ASTer) сконвертитьСтрокувПараметры(in string) (ret Параметры) { params := strings.Split(in, ",") for _, p := range params { ps := strings.Fields(p) параметр := Параметр{} switch len(ps) { case 0: continue case 1: параметр = append(параметр, ps[0]) case 2: параметр = append(параметр, ps[0], ps[1]) } ret = append(ret, параметр) } return } func (a *ASTer) перезаписатьInitФайл() error { err := переименоватьФайлСоВременем(INIT_TEST_GO_FNAME) if err != nil { return errors.Join(err, errors.New("cant backup orig file")) } f, err := os.Create(INIT_TEST_GO_FNAME) if err != nil { return errors.Join(err, errors.New("cant rewrite orig file")) } defer f.Close() return format.Node(f, a.init_test_fset, a.init_test_node) } func (a *ASTer) перезаписатьФайл() error { err := переименоватьФайлСоВременем(a.pkg_test_go_fname) if err != nil { return errors.Join(err, errors.New("cant backup orig file")) } f, err := os.Create(a.pkg_test_go_fname) if err != nil { return errors.Join(err, errors.New("cant rewrite orig file")) } defer f.Close() return format.Node(f, a.pkg_test_fset, a.pkg_test_node) }