Test code extraction for exercises with separate file for test cases

This commit also support arbitrary test data variable name when test function
includes multiple assignments.

Author: meoconbatu

go.mod

@@ -7,4 +7,4 @@ require (
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/stretchr/testify v1.8.2 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
-)
+)

go.sum

@@ -13,4 +13,4 @@ github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
-gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

testrunner/ast.go

@@ -3,12 +3,14 @@ package testrunner
 import (
 	"bytes"
 	"errors"
+	"fmt"
 	"go/ast"
 	"go/format"
 	"go/parser"
 	"go/printer"
 	"go/token"
 	"log"
+	"path/filepath"
 	"regexp"
 	"strconv"
 	"strings"
@@ -34,6 +36,7 @@ type rootLevelTest struct {
 	fileName string
 	code     string
 	taskID   uint64
+	pkgName  string
 }
 
 // FindAllRootLevelTests parses the test file and extracts the name,
@@ -60,6 +63,7 @@ func FindAllRootLevelTests(fileName string) []rootLevelTest {
 				fileName: fileName,
 				code:     buf.String(),
 				taskID:   taskID,
+				pkgName:  file.Name.Name,
 			})
 		}
 	}
@@ -95,16 +99,19 @@ func findTaskID(doc *ast.CommentGroup) uint64 {
 }
 
 // generate simplified test code corresponding to a subtest
-func getSubCode(test string, sub string, code string, file string) string {
+func getSubCode(test string, sub string, code string, file string, pkgName string) string {
+	pkgLine := fmt.Sprintf("package %s\n", pkgName)
 	fset := token.NewFileSet()
 	f, err := parser.ParseFile(
-		fset, file, "package main\n"+code, parser.ParseComments,
+		fset, file, pkgLine+code, parser.ParseComments,
 	)
 	if err != nil {
 		log.Printf("warning: '%s' not parsed from '%s': %s", test, file, err)
 		return ""
 	}
 
+	resolveTestData(fset, f, file)
+
 	fAST, ok := f.Decls[0].(*ast.FuncDecl)
 	if !ok {
 		log.Println("warning: first subtest declaration must be a function")
@@ -113,7 +120,7 @@ func getSubCode(test string, sub string, code string, file string) string {
 
 	fbAST := fAST.Body.List // f.Decls[0].Body.List
 
-	astInfo, err := findTestDataAndRange(fbAST)
+	astInfo, err := findTestDataAndRange(fbAST, fset)
 	if err != nil {
 		log.Printf("warning: could not find test table and/or range: %v\n", err)
 		return ""
@@ -146,36 +153,33 @@ func getSubCode(test string, sub string, code string, file string) string {
 		log.Println("warning: failed to format extracted AST for subtest")
 		return ""
 	}
-	return strings.TrimSpace(strings.TrimPrefix(buf.String(), "package main"))
+	if astInfo.testDataAstIdx != -1 { // testDataAst is already in the test function
+		return strings.TrimSpace(strings.TrimPrefix(buf.String(), pkgLine))
+	}
+	return insertTestDataASTIntoFunc(fset, astInfo.testDataAst, fAST.Body, buf.Bytes(), pkgLine)
 }
 
-func findTestDataAndRange(stmtList []ast.Stmt) (subTestAstInfo, error) {
+func findTestDataAndRange(stmtList []ast.Stmt, fset *token.FileSet) (subTestAstInfo, error) {
 	result := subTestAstInfo{}
-
+	posToIndex := make(map[token.Position]int)
 	for i := range stmtList {
-		assignCandidate, ok := stmtList[i].(*ast.AssignStmt)
-		if ok && result.testDataAst == nil {
-			result.testDataAst = assignCandidate
-			result.testDataAstIdx = i
-		} else if ok {
-			identifier, isIdentifier := assignCandidate.Lhs[0].(*ast.Ident)
-			if !isIdentifier {
-				continue
-			}
-			// Overwrite the assignment we already found in case there is an
-			// assignment to a "tests" variable.
-			if identifier.Name == "tests" {
+		posToIndex[fset.Position(stmtList[i].Pos())] = i
+		if rangeCandidate, ok := stmtList[i].(*ast.RangeStmt); ok {
+			assignCandidate := getTestDataAssignFromRange(rangeCandidate)
+			if assignCandidate != nil {
+				// check if assignCandidate is in the same function with rangeCandidate
+				if idx, ok := posToIndex[fset.Position(assignCandidate.Pos())]; ok &&
+					fset.File(assignCandidate.Pos()).Name() == fset.File(rangeCandidate.Pos()).Name() {
+					result.testDataAstIdx = idx
+				} else {
+					result.testDataAstIdx = -1
+				}
 				result.testDataAst = assignCandidate
-				result.testDataAstIdx = i
+				result.rangeAst = rangeCandidate
+				result.rangeAstIdx = i
+				return result, nil
 			}
-		}
-
-		rangeCandidate, ok := stmtList[i].(*ast.RangeStmt)
-		// If we found a range after we already found an assignment, we are good to go.
-		if ok && result.testDataAst != nil {
-			result.rangeAst = rangeCandidate
-			result.rangeAstIdx = i
-			return result, nil
+			return subTestAstInfo{}, errors.New("failed to find assignment in sub-test")
 		}
 	}
 
@@ -185,6 +189,24 @@ func findTestDataAndRange(stmtList []ast.Stmt) (subTestAstInfo, error) {
 
 	return subTestAstInfo{}, errors.New("failed to find range statement in sub-test")
 }
+func getTestDataAssignFromRange(rangeAst *ast.RangeStmt) *ast.AssignStmt {
+	spec := rangeAst.X.(*ast.Ident).Obj.Decl
+	if assignStmt, ok := spec.(*ast.AssignStmt); ok {
+		return assignStmt
+	}
+	if valueSpec, ok := spec.(*ast.ValueSpec); ok {
+		lhs := make([]ast.Expr, len(valueSpec.Names))
+		for i, name := range valueSpec.Names {
+			lhs[i] = name
+		}
+		return &ast.AssignStmt{
+			Lhs: lhs,
+			Tok: token.DEFINE,
+			Rhs: valueSpec.Values,
+		}
+	}
+	return nil
+}
 
 // validate the test data assignment and return the associated metadata
 func processTestDataAssgn(sub string, assgn *ast.AssignStmt) (*subTData, bool) {
@@ -290,3 +312,44 @@ func processRange(metadata *subTData, rastmt *ast.RangeStmt) bool {
 	metadata.subTest = body
 	return true
 }
+
+// resolveTestData resolves test data variable declared in cases_test.go (if exists)
+func resolveTestData(fset *token.FileSet, f *ast.File, file string) {
+	filedata := filepath.Join(filepath.Dir(file), "cases_test.go")
+	fdata, _ := parser.ParseFile(fset, filedata, nil, parser.ParseComments)
+
+	// NewPackage func always return errors because f files's missing import part
+	// so ignore checking the returned errors
+	if fdata != nil {
+		_, _ = ast.NewPackage(fset, map[string]*ast.File{file: f, filedata: fdata}, nil, nil)
+	} else {
+		_, _ = ast.NewPackage(fset, map[string]*ast.File{file: f}, nil, nil)
+	}
+}
+
+// insertTestDataASTIntoFunc inserts testDataAst into the first line of fbAST function's body
+func insertTestDataASTIntoFunc(fset *token.FileSet, testDataAst *ast.AssignStmt, fbAST *ast.BlockStmt, fileText []byte, pkgLine string) string {
+	buf := bytes.Buffer{}
+
+	p := fset.Position(fbAST.Lbrace).Offset + 1
+
+	// write the beginning of fileText to func (...) {
+	buf.Write(fileText[:p+1])
+
+	// write test data assign stmt
+	if err := format.Node(&buf, fset, testDataAst); err != nil {
+		log.Println("warning: failed to format extracted AST for subtest")
+		return ""
+	}
+	// write the rest of fileText
+	buf.Write(fileText[p+1:])
+
+	// because assign stmt is extracted from different file, its indentation is different from fileText
+	// so need to reformat
+	src, err := format.Source((buf.Bytes()))
+	if err != nil {
+		log.Println("warning: failed to format extracted AST for subtest")
+		return ""
+	}
+	return strings.TrimSpace(strings.TrimPrefix(string(src), pkgLine))
+}

testrunner/extract.go

@@ -53,7 +53,7 @@ func ExtractTestCodeAndTaskID(rootLevelTests map[string]rootLevelTest, testName
 		return rootLevelTest.code, rootLevelTest.taskID
 	}
 	defer handleASTPanic()
-	subtc := getSubCode(test, subtest, rootLevelTest.code, rootLevelTest.fileName)
+	subtc := getSubCode(test, subtest, rootLevelTest.code, rootLevelTest.fileName, rootLevelTest.pkgName)
 	if len(subtc) == 0 {
 		return rootLevelTest.code, rootLevelTest.taskID
 	}

testrunner/extract_test.go

@@ -191,6 +191,8 @@ func TestExtractTestCode(t *testing.T) {
   }`,
 		},
 	}
+	tests = append(tests, testsDataSeparate...)
+	tests = append(tests, testsMultiAssignStmt...)
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			code, _ := ExtractTestCodeAndTaskID(rootLevelTestsMap, tt.testName)
@@ -221,3 +223,108 @@ func TestExtractTestCode(t *testing.T) {
 		})
 	}
 }
+
+var testsDataSeparate = []struct {
+	name     string
+	testName string
+	testFile string
+	code     string
+}{
+	{
+		name:     "working subtest with separate test data",
+		testName: "TestParseCard_Separate/parse_jack",
+		testFile: filepath.Join("testdata", "concept", "conditionals", "conditionals_test.go"),
+		code: `func TestParseCard_Separate(t *testing.T) {
+	tt := struct {
+		name string
+		card string
+		want int
+	}{
+		name: "parse jack",
+		card: "jack",
+		want: 10,
+	}
+	
+	if got := ParseCard(tt.card); got != tt.want {
+		t.Errorf("ParseCard(%s) = %d, want %d", tt.card, got, tt.want)
+	}
+
+}`,
+	}, {
+		name:     "missing / not found subtest with separate test data",
+		testName: "TestParseCard_Separate/parse_missing_subtests",
+		testFile: filepath.Join("testdata", "concept", "conditionals", "conditionals_test.go"),
+		code: `func TestParseCard_Separate(t *testing.T) {
+	for _, tt := range testcases {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := ParseCard(tt.card); got != tt.want {
+				t.Errorf("ParseCard(%s) = %d, want %d", tt.card, got, tt.want)
+			}
+		})
+	}
+}`,
+	}, {
+		name:     "multiple statements with separate test data",
+		testName: "TestBlackjack_Separate/blackjack_with_ten_(ace_first)",
+		testFile: filepath.Join("testdata", "concept", "conditionals", "conditionals_test.go"),
+		code: `func TestBlackjack_Separate(t *testing.T) {
+	tt := struct {
+		name string
+		hand hand
+		want bool
+	}{
+		name: "blackjack with ten (ace first)",
+		hand: hand{card1: "ace", card2: "ten"},
+		want: true,
+	}
+	someAssignment := "test"
+	fmt.Println(someAssignment)
+
+	_ = "literally anything"
+
+	got := IsBlackjack(tt.hand.card1, tt.hand.card2)
+	if got != tt.want {
+		t.Errorf("IsBlackjack(%s, %s) = %t, want %t", tt.hand.card1, tt.hand.card2, got, tt.want)
+	}
+
+	// Additional statements should be included
+	fmt.Println("the whole block")
+	fmt.Println("should be returned")
+}`,
+	},
+}
+var testsMultiAssignStmt = []struct {
+	name     string
+	testName string
+	testFile string
+	code     string
+}{
+	{
+		name:     "subtest with arbitrary test data variable name, additional assign statements above and below test data",
+		testName: "TestSubtest_MultiAssignStmt/parse_king",
+		testFile: filepath.Join("testdata", "concept", "conditionals", "conditionals_test.go"),
+		code: `func TestSubtest_MultiAssignStmt(t *testing.T) {
+	someAssignment := "test"
+
+	tt := struct {
+		name string
+		card string
+		want int
+	}{
+		name: "parse king",
+		card: "king",
+		want: 10,
+	}
+
+	someAssignment2 := "test2"
+
+	if got := ParseCard(tt.card); got != tt.want {
+		t.Errorf("ParseCard(%s) = %d, want %d", tt.card, got, tt.want)
+	}
+
+	// Additional statements should be included
+	fmt.Println("the whole block")
+	fmt.Println("should be returned")
+}`,
+	},
+}

testrunner/testdata/concept/conditionals/cases_test.go

@@ -19,3 +19,62 @@ var allergicToTests = []struct {
 		expected: false,
 	},
 }
+var testcases = []struct {
+	name string
+	card string
+	want int
+}{
+	{
+		name: "parse two",
+		card: "two",
+		want: 2,
+	},
+	{
+		name: "parse jack",
+		card: "jack",
+		want: 10,
+	},
+	{
+		name: "parse king",
+		card: "king",
+		want: 10,
+	},
+}
+
+type hand struct {
+	card1, card2 string
+}
+
+var testcases2 = []struct {
+	name string
+	hand hand
+	want bool
+}{
+	{
+		name: "blackjack with ten (ace first)",
+		hand: hand{card1: "ace", card2: "ten"},
+		want: true,
+	},
+	{
+		name: "blackjack with jack (ace first)",
+		hand: hand{card1: "ace", card2: "jack"},
+		want: true,
+	},
+	{
+		name: "blackjack with queen (ace first)",
+		hand: hand{
+			card1: "ace", card2: "queen"
+		},
+		want: true,
+	},
+	{
+		name: "blackjack with king (ace first)",
+		hand: hand{card1: "ace", card2: "king"},
+		want: true,
+	},
+	{
+		name: "no blackjack with eight and five",
+		hand: hand{card2: "eight", card1: "five"},
+		want: false,
+	},
+}