Add model generator based on AST

pkg/codegen/modelgenerate/generate_testfile.go

@@ -0,0 +1,251 @@
+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")},
+							},
+						}},
+					},
+				},
+			},
+		},
+	}
+}