internal/godoc: don't remove test file nodes

It's possible that a test file in the main package (as opposed to a
_test package) can have a playable example. So don't remove any
AST nodes from test files.

With this change, FetchModule produces the same doc HTML with and
without the RemoveUnusedAST and InsertPackageSource experiments for
all modules with the top 5000 packages.

Change-Id: Ib003ef5d2fd1f28694953fd414178a1bb8eff965
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/260678
Trust: Jonathan Amsterdam <jba@google.com>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
TryBot-Result: kokoro <noreply+kokoro@google.com>
Reviewed-by: Julie Qiu <julie@golang.org>
diff --git a/internal/godoc/godoc.go b/internal/godoc/godoc.go
index eb66d43..35272a0 100644
--- a/internal/godoc/godoc.go
+++ b/internal/godoc/godoc.go
@@ -60,11 +60,14 @@
 // AddFile adds a file to the Package. After it returns, the contents of the ast.File
 // are unsuitable for anything other than the methods of this package.
 func (p *Package) AddFile(f *ast.File, removeNodes bool) {
-	if removeNodes {
+	filename := p.Fset.Position(f.Package).Filename
+	// Don't trim anything from a test file or one in a XXX_test package; it
+	// may be part of a playable example.
+	if removeNodes && !strings.HasSuffix(filename, "_test.go") && !strings.HasSuffix(f.Name.Name, "_test") {
 		removeUnusedASTNodes(f)
 	}
 	p.Files = append(p.Files, &File{
-		Name: p.Fset.Position(f.Package).Filename,
+		Name: filename,
 		AST:  f,
 	})
 }
@@ -72,11 +75,6 @@
 // removeUnusedASTNodes removes parts of the AST not needed for documentation.
 // It doesn't remove unexported consts, vars or types, although it probably could.
 func removeUnusedASTNodes(pf *ast.File) {
-	// Don't trim anything from a file in a XXX_test package; it
-	// may be part of a playable example.
-	if strings.HasSuffix(pf.Name.Name, "_test") {
-		return
-	}
 	var decls []ast.Decl
 	for _, d := range pf.Decls {
 		if f, ok := d.(*ast.FuncDecl); ok {