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

package modelgenerate

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

	"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")},
				},
			},
		},
	}

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

	return &ast.File{
		Name: ast.NewIdent(testpackageName),
		Decls: []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("db"),
					},
					&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"`}},
				},
			},
			// var TestDB *DB
			testDBDecl,

			// 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,

			&ast.FuncDecl{
				Name: ast.NewIdent("TestCreateUpdateDelete" + tbl.TypeName),
				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.TypeName),
								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.TypeName),
							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.TypeName + "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.TypeName),
							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.TypeName + "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.TypeName),
							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.TypeName + "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.SelectorExpr{X: ast.NewIdent("db"), Sel: ast.NewIdent("ErrNotInDB")},
							},
						}},
					},
				},
			},
		},
	}
}