diff --git a/pkg/codegen/modelgenerate/generate_testfile.go b/pkg/codegen/modelgenerate/generate_testfile.go index a3ca2f0..904897f 100644 --- a/pkg/codegen/modelgenerate/generate_testfile.go +++ b/pkg/codegen/modelgenerate/generate_testfile.go @@ -101,6 +101,22 @@ func GenerateModelTestAST(tbl pkgschema.Table, schema pkgschema.Schema, gomodNam hasCreatedAt, hasUpdatedAt := tbl.HasAutoTimestamps() updateTestFields := UpdateTestFields(tbl) + uniqueColumns := map[string]bool{} + for _, index := range schema.Indexes { + if index.TableName != tbl.TableName || !index.IsUnique { + continue + } + for _, columnName := range index.Columns { + uniqueColumns[columnName] = true + } + } + hasRandomUniqueField := false + for _, c := range updateTestFields { + if uniqueColumns[c.Name] { + hasRandomUniqueField = true + break + } + } requiredForeignKeys := []pkgschema.Column{} for _, c := range tbl.Columns { // Only ordinary object foreign keys use a distinct TargetObjectID Go type; @@ -141,31 +157,71 @@ func GenerateModelTestAST(tbl pkgschema.Table, schema pkgschema.Schema, gomodNam }, }, Body: &ast.BlockStmt{ - List: []ast.Stmt{ - &ast.ReturnStmt{ - Results: []ast.Expr{ - &ast.CompositeLit{ - Type: ast.NewIdent(tbl.GoTypeName), - Elts: func() (ret []ast.Expr) { - // `UpdateTestFields()` skips foreign keys, so this will not cause duplicates - for i, c := range requiredForeignKeys { - ret = append(ret, &ast.KeyValueExpr{ - Key: ast.NewIdent(c.GoFieldName()), - Value: ast.NewIdent(requiredForeignKeyParamNames[i]), - }) - } - for _, c := range updateTestFields { - ret = append(ret, &ast.KeyValueExpr{ - Key: ast.NewIdent(c.GoFieldName()), - Value: SampleValue(c, 0), - }) - } - return - }(), + List: append( + func() []ast.Stmt { + if !hasRandomUniqueField { + return nil + } + return []ast.Stmt{&ast.AssignStmt{ + Lhs: []ast.Expr{ast.NewIdent("id")}, + Tok: token.DEFINE, + Rhs: []ast.Expr{&ast.CallExpr{Fun: &ast.SelectorExpr{X: ast.NewIdent("rand"), Sel: ast.NewIdent("Int")}}}, + }} + }(), + []ast.Stmt{ + &ast.ReturnStmt{ + Results: []ast.Expr{ + &ast.CompositeLit{ + Type: ast.NewIdent(tbl.GoTypeName), + Elts: func() (ret []ast.Expr) { + // `UpdateTestFields()` skips foreign keys, so this will not cause duplicates + for i, c := range requiredForeignKeys { + ret = append(ret, &ast.KeyValueExpr{ + Key: ast.NewIdent(c.GoFieldName()), + Value: ast.NewIdent(requiredForeignKeyParamNames[i]), + }) + } + for _, c := range updateTestFields { + value := SampleValue(c, 0) + if uniqueColumns[c.Name] { + switch c.Type { + case "integer", "int": + value = ast.NewIdent("id") + case "text": + value = &ast.CallExpr{ + Fun: &ast.SelectorExpr{X: ast.NewIdent("fmt"), Sel: ast.NewIdent("Sprintf")}, + Args: []ast.Expr{ + &ast.BasicLit{Kind: token.STRING, Value: `"%s-%d"`}, + SampleValue(c, 0), + ast.NewIdent("id"), + }, + } + case "blob": + value = &ast.CallExpr{ + Fun: &ast.ArrayType{Elt: ast.NewIdent("byte")}, + Args: []ast.Expr{&ast.CallExpr{ + Fun: &ast.SelectorExpr{X: ast.NewIdent("fmt"), Sel: ast.NewIdent("Sprintf")}, + Args: []ast.Expr{ + &ast.BasicLit{Kind: token.STRING, Value: `"%s-%d"`}, + &ast.CallExpr{Fun: ast.NewIdent("string"), Args: []ast.Expr{SampleValue(c, 0)}}, + ast.NewIdent("id"), + }, + }}, + } + } + } + ret = append(ret, &ast.KeyValueExpr{ + Key: ast.NewIdent(c.GoFieldName()), + Value: value, + }) + } + return + }(), + }, }, }, - }, - }, + }..., + ), }, } @@ -595,23 +651,33 @@ func GenerateModelTestAST(tbl pkgschema.Table, schema pkgschema.Schema, gomodNam Decls: append([]ast.Decl{ &ast.GenDecl{ Tok: token.IMPORT, - Specs: []ast.Spec{ - &ast.ImportSpec{Path: &ast.BasicLit{Kind: token.STRING, Value: `"testing"`}}, - &ast.ImportSpec{ - Path: &ast.BasicLit{Kind: token.STRING, Value: `"git.offline-twitter.com/offline-labs/gas-stack/pkg/db"`}, - Name: ast.NewIdent("."), - }, - &ast.ImportSpec{ - Path: &ast.BasicLit{Kind: token.STRING, Value: `"git.offline-twitter.com/offline-labs/gas-stack/pkg/must"`}, - }, - &ast.ImportSpec{ - Path: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf(`"%s/pkg/%s"`, gomodName, packageName)}, - Name: ast.NewIdent("."), - }, - &ast.ImportSpec{Path: &ast.BasicLit{Kind: token.STRING, Value: `"github.com/go-test/deep"`}}, - &ast.ImportSpec{Path: &ast.BasicLit{Kind: token.STRING, Value: `"github.com/stretchr/testify/assert"`}}, - &ast.ImportSpec{Path: &ast.BasicLit{Kind: token.STRING, Value: `"github.com/stretchr/testify/require"`}}, - }, + Specs: func() []ast.Spec { + specs := []ast.Spec{} + if hasRandomUniqueField { + specs = append(specs, + &ast.ImportSpec{Path: &ast.BasicLit{Kind: token.STRING, Value: `"fmt"`}}, + &ast.ImportSpec{Path: &ast.BasicLit{Kind: token.STRING, Value: `"math/rand"`}}, + ) + } + specs = append(specs, []ast.Spec{ + &ast.ImportSpec{Path: &ast.BasicLit{Kind: token.STRING, Value: `"testing"`}}, + &ast.ImportSpec{ + Path: &ast.BasicLit{Kind: token.STRING, Value: `"git.offline-twitter.com/offline-labs/gas-stack/pkg/db"`}, + Name: ast.NewIdent("."), + }, + &ast.ImportSpec{ + Path: &ast.BasicLit{Kind: token.STRING, Value: `"git.offline-twitter.com/offline-labs/gas-stack/pkg/must"`}, + }, + &ast.ImportSpec{ + Path: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf(`"%s/pkg/%s"`, gomodName, packageName)}, + Name: ast.NewIdent("."), + }, + &ast.ImportSpec{Path: &ast.BasicLit{Kind: token.STRING, Value: `"github.com/go-test/deep"`}}, + &ast.ImportSpec{Path: &ast.BasicLit{Kind: token.STRING, Value: `"github.com/stretchr/testify/assert"`}}, + &ast.ImportSpec{Path: &ast.BasicLit{Kind: token.STRING, Value: `"github.com/stretchr/testify/require"`}}, + }...) + return specs + }(), }, }, testList...), }