gas-stack/pkg/codegen/modelgenerate/generate_testfile.go

package modelgenerate

import (
	"fmt"
	"go/ast"
	"go/token"

	"github.com/jinzhu/inflection"

	"git.offline-twitter.com/offline-labs/gas-stack/pkg/schema"
)

// GenerateModelTestAST produces an AST for a starter test file for a given model.
func GenerateModelTestAST(tbl schema.Table, gomodName string) *ast.File {
	packageName := "db"
	testpackageName := packageName + "_test"

	testDBDecl := &ast.GenDecl{
		Tok: token.VAR,
		Specs: []ast.Spec{
			&ast.ValueSpec{
				Names: []*ast.Ident{ast.NewIdent("TestDB")},
				Type:  &ast.StarExpr{X: ast.NewIdent("DB")},
			},
		},
	}

	initFuncDecl := &ast.FuncDecl{
		Name: ast.NewIdent("init"),
		Type: &ast.FuncType{Params: &ast.FieldList{}},
		Body: &ast.BlockStmt{
			List: []ast.Stmt{
				&ast.AssignStmt{
					Lhs: []ast.Expr{ast.NewIdent("TestDB")},
					Tok: token.ASSIGN,
					Rhs: []ast.Expr{&ast.CallExpr{
						Fun:  ast.NewIdent("MakeDB"),
						Args: []ast.Expr{&ast.BasicLit{Kind: token.STRING, Value: `"tmp"`}},
					}},
				},
			},
		},
	}

	makeDBHelperDecl := &ast.FuncDecl{
		Name: ast.NewIdent("MakeDB"),
		Type: &ast.FuncType{
			Params: &ast.FieldList{
				List: []*ast.Field{{
					Names: []*ast.Ident{ast.NewIdent("dbName")},
					Type:  ast.NewIdent("string"),
				}},
			},
			Results: &ast.FieldList{
				List: []*ast.Field{{Type: &ast.StarExpr{X: ast.NewIdent("DB")}}},
			},
		},
		Body: &ast.BlockStmt{
			List: []ast.Stmt{
				// db := Must(Create(fmt.Sprintf("file:%s?mode=memory&cache=shared", dbName)))
				&ast.AssignStmt{
					Lhs: []ast.Expr{ast.NewIdent("db")},
					Tok: token.DEFINE,
					Rhs: []ast.Expr{&ast.CallExpr{
						Fun: ast.NewIdent("Must"),
						Args: []ast.Expr{&ast.CallExpr{
							Fun: ast.NewIdent("Create"),
							Args: []ast.Expr{&ast.CallExpr{
								Fun: ast.NewIdent("fmt.Sprintf"),
								Args: []ast.Expr{
									&ast.BasicLit{Kind: token.STRING, Value: `"file:%s?mode=memory&cache=shared"`},
									ast.NewIdent("dbName"),
								},
							}},
						}},
					}},
				},
				// return db
				&ast.ReturnStmt{
					Results: []ast.Expr{ast.NewIdent("db")},
				},
			},
		},
	}

	// func MakeItem() Item { return Item{} }
	makeItemFunc := &ast.FuncDecl{
		Name: ast.NewIdent("Make" + tbl.GoTypeName),
		Type: &ast.FuncType{
			Params: &ast.FieldList{},
			Results: &ast.FieldList{
				List: []*ast.Field{
					{Type: ast.NewIdent(tbl.GoTypeName)},
				},
			},
		},
		Body: &ast.BlockStmt{
			List: []ast.Stmt{
				&ast.ReturnStmt{
					Results: []ast.Expr{
						&ast.CompositeLit{
							Type: ast.NewIdent(tbl.GoTypeName),
						},
					},
				},
			},
		},
	}

	testObj := ast.NewIdent("item")
	testObj2 := ast.NewIdent("item2")
	fieldName := ast.NewIdent("Description")
	description1 := `"an item"`
	description2 := `"a big item"`

	testCreateUpdateDelete := &ast.FuncDecl{
		Name: ast.NewIdent("TestCreateUpdateDelete" + tbl.GoTypeName),
		Type: &ast.FuncType{
			Params: &ast.FieldList{
				List: []*ast.Field{{
					Names: []*ast.Ident{ast.NewIdent("t")},
					Type:  ast.NewIdent("*testing.T"),
				}},
			},
		},
		Body: &ast.BlockStmt{
			List: []ast.Stmt{
				// item := Item{Description: "an item"}
				&ast.AssignStmt{
					Lhs: []ast.Expr{testObj},
					Tok: token.DEFINE,
					Rhs: []ast.Expr{&ast.CompositeLit{
						Type: ast.NewIdent(tbl.GoTypeName),
						Elts: []ast.Expr{
							&ast.KeyValueExpr{
								Key:   fieldName,
								Value: &ast.BasicLit{Kind: token.STRING, Value: description1},
							},
						},
					}},
				},

				// TestDB.SaveItem(&item)
				&ast.ExprStmt{X: &ast.CallExpr{
					Fun:  ast.NewIdent("TestDB.Save" + tbl.GoTypeName),
					Args: []ast.Expr{&ast.UnaryExpr{Op: token.AND, X: testObj}},
				}},

				// require.NotZero(t, item.ID)
				&ast.ExprStmt{X: &ast.CallExpr{
					Fun:  ast.NewIdent("require.NotZero"),
					Args: []ast.Expr{ast.NewIdent("t"), &ast.SelectorExpr{X: testObj, Sel: ast.NewIdent("ID")}},
				}},

				// item2 := Must(TestDB.GetItemByID(item.ID))
				&ast.AssignStmt{
					Lhs: []ast.Expr{testObj2},
					Tok: token.DEFINE,
					Rhs: []ast.Expr{&ast.CallExpr{
						Fun: ast.NewIdent("Must"),
						Args: []ast.Expr{&ast.CallExpr{
							Fun:  ast.NewIdent("TestDB.Get" + tbl.GoTypeName + "ByID"),
							Args: []ast.Expr{&ast.SelectorExpr{X: testObj, Sel: ast.NewIdent("ID")}},
						}},
					}},
				},

				// assert.Equal(t, "an item", item2.Description)
				&ast.ExprStmt{X: &ast.CallExpr{
					Fun: ast.NewIdent("assert.Equal"),
					Args: []ast.Expr{
						ast.NewIdent("t"),
						&ast.BasicLit{Kind: token.STRING, Value: description1},
						&ast.SelectorExpr{X: testObj2, Sel: fieldName},
					},
				}},

				// item.Description = "a big item"
				&ast.AssignStmt{
					Lhs: []ast.Expr{&ast.SelectorExpr{X: testObj, Sel: fieldName}},
					Tok: token.ASSIGN,
					Rhs: []ast.Expr{&ast.BasicLit{Kind: token.STRING, Value: description2}},
				},

				// TestDB.SaveItem(&item)
				&ast.ExprStmt{X: &ast.CallExpr{
					Fun:  ast.NewIdent("TestDB.Save" + tbl.GoTypeName),
					Args: []ast.Expr{&ast.UnaryExpr{Op: token.AND, X: testObj}},
				}},

				// item2 = Must(TestDB.GetItemByID(item.ID))
				&ast.AssignStmt{
					Lhs: []ast.Expr{testObj2},
					Tok: token.ASSIGN,
					Rhs: []ast.Expr{&ast.CallExpr{
						Fun: ast.NewIdent("Must"),
						Args: []ast.Expr{&ast.CallExpr{
							Fun:  ast.NewIdent("TestDB.Get" + tbl.GoTypeName + "ByID"),
							Args: []ast.Expr{&ast.SelectorExpr{X: testObj, Sel: ast.NewIdent("ID")}},
						}},
					}},
				},

				// assert.Equal(t, item.Description, item2.Description)
				&ast.ExprStmt{X: &ast.CallExpr{
					Fun: ast.NewIdent("assert.Equal"),
					Args: []ast.Expr{
						ast.NewIdent("t"),
						&ast.SelectorExpr{X: testObj, Sel: fieldName},
						&ast.SelectorExpr{X: testObj2, Sel: fieldName},
					},
				}},

				// TestDB.DeleteItem(item)
				&ast.ExprStmt{X: &ast.CallExpr{
					Fun:  ast.NewIdent("TestDB.Delete" + tbl.GoTypeName),
					Args: []ast.Expr{testObj},
				}},

				// _, err := TestDB.GetItemByID(item.ID)
				&ast.AssignStmt{
					Lhs: []ast.Expr{ast.NewIdent("_"), ast.NewIdent("err")},
					Tok: token.DEFINE,
					Rhs: []ast.Expr{&ast.CallExpr{
						Fun:  ast.NewIdent("TestDB.Get" + tbl.GoTypeName + "ByID"),
						Args: []ast.Expr{&ast.SelectorExpr{X: testObj, Sel: ast.NewIdent("ID")}},
					}},
				},

				// assert.ErrorIs(t, err, db.ErrNotInDB)
				&ast.ExprStmt{X: &ast.CallExpr{
					Fun: ast.NewIdent("assert.ErrorIs"),
					Args: []ast.Expr{
						ast.NewIdent("t"),
						ast.NewIdent("err"),
						ast.NewIdent("ErrNotInDB"),
					},
				}},
			},
		},
	}

	testGetAll := &ast.FuncDecl{
		Name: ast.NewIdent("TestGetAll" + inflection.Plural(tbl.GoTypeName)),
		Type: &ast.FuncType{Params: &ast.FieldList{List: []*ast.Field{
			{Names: []*ast.Ident{ast.NewIdent("t")}, Type: &ast.StarExpr{X: ast.NewIdent("testing.T")}},
		}}, Results: nil},
		Body: &ast.BlockStmt{
			List: []ast.Stmt{
				&ast.AssignStmt{
					Lhs: []ast.Expr{ast.NewIdent("_")},
					Tok: token.ASSIGN,
					Rhs: []ast.Expr{&ast.CallExpr{
						Fun: &ast.SelectorExpr{
							X:   ast.NewIdent("TestDB"),
							Sel: ast.NewIdent("GetAll" + inflection.Plural(tbl.GoTypeName)),
						},
					}},
				},
			},
		},
	}

	shouldIncludeTestFkCheck := false
	testFkChecking := &ast.FuncDecl{
		Name: ast.NewIdent("Test" + tbl.GoTypeName + "FkChecking"),
		Type: &ast.FuncType{
			Params: &ast.FieldList{
				List: []*ast.Field{
					{
						Names: []*ast.Ident{ast.NewIdent("t")},
						Type: &ast.StarExpr{
							X: &ast.SelectorExpr{
								X:   ast.NewIdent("testing"),
								Sel: ast.NewIdent("T"),
							},
						},
					},
				},
			},
		},
		Body: &ast.BlockStmt{
			List: func() []ast.Stmt {
				// post := MakePost()
				stmts := []ast.Stmt{
					&ast.AssignStmt{
						Lhs: []ast.Expr{ast.NewIdent(tbl.VarName)},
						Tok: token.DEFINE,
						Rhs: []ast.Expr{
							&ast.CallExpr{
								Fun: ast.NewIdent("Make" + tbl.GoTypeName),
							},
						},
					},
				}
				for _, col := range tbl.Columns {
					if col.IsForeignKey {
						shouldIncludeTestFkCheck = true
						stmts = append(stmts, []ast.Stmt{

							// post.QuotedPostID = 94354538969386985
							&ast.AssignStmt{
								Lhs: []ast.Expr{
									&ast.SelectorExpr{
										X:   ast.NewIdent(tbl.VarName),
										Sel: ast.NewIdent(fkFieldName(col)),
									},
								},
								Tok: token.ASSIGN,
								Rhs: []ast.Expr{
									&ast.BasicLit{
										Kind:  token.INT,
										Value: "94354538969386985",
									},
								},
							},

							// err := db.SavePost(&post)
							&ast.AssignStmt{
								Lhs: []ast.Expr{ast.NewIdent("err")},
								Tok: token.DEFINE,
								Rhs: []ast.Expr{
									&ast.CallExpr{
										Fun: &ast.SelectorExpr{
											X:   ast.NewIdent("TestDB"),
											Sel: ast.NewIdent("Save" + tbl.GoTypeName),
										},
										Args: []ast.Expr{
											&ast.UnaryExpr{
												Op: token.AND,
												X:  ast.NewIdent(tbl.VarName),
											},
										},
									},
								},
							},

							// assertForeignKeyError(t, err, "QuotedPostID", post.QuotedPostID)
							&ast.ExprStmt{
								X: &ast.CallExpr{
									Fun: ast.NewIdent("AssertForeignKeyError"),
									Args: []ast.Expr{
										ast.NewIdent("t"),
										ast.NewIdent("err"),
										&ast.BasicLit{
											Kind:  token.STRING,
											Value: fmt.Sprintf("%q", fkFieldName(col)),
										},
										&ast.SelectorExpr{
											X:   ast.NewIdent(tbl.VarName),
											Sel: ast.NewIdent(fkFieldName(col)),
										},
									},
								},
							},
						}...)
					}
				}
				return stmts
			}(),
		},
	}

	testList := []ast.Decl{
		// var TestDB *DB
		testDBDecl,
		makeItemFunc,

		// func init() { TestDB = MakeDB("tmp") }
		initFuncDecl,

		// func MakeDB(dbName string) *DB { db := Must(Create(fmt.Sprintf("file:%s?mode=memory&cache=shared", dbName))); return db }
		makeDBHelperDecl,

		testCreateUpdateDelete,
		testGetAll,
	}
	if shouldIncludeTestFkCheck {
		testList = append(testList, testFkChecking)
	}
	return &ast.File{
		Name: ast.NewIdent(testpackageName),
		Decls: append([]ast.Decl{
			&ast.GenDecl{
				Tok: token.IMPORT,
				Specs: []ast.Spec{
					&ast.ImportSpec{Path: &ast.BasicLit{Kind: token.STRING, Value: `"fmt"`}},
					&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/flowutils"`},
						Name: ast.NewIdent("."),
					},
					&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/stretchr/testify/assert"`}},
					&ast.ImportSpec{Path: &ast.BasicLit{Kind: token.STRING, Value: `"github.com/stretchr/testify/require"`}},
				},
			},
		}, testList...),
	}
}