internal/fetch/internal/doc: fix detection of whole file examples

This CL uses the fix from CL 222477 (golang/go#38409), to fix detection
of whole file examples.

Fixes golang/go#39006
Fixes golang/go#39727

Change-Id: Iaf9240f7637a1723e681cdbb0d06bf36a43d8d2a
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/239180
Reviewed-by: Jonathan Amsterdam <[email protected]>
Reviewed-by: Dmitri Shuralyov <[email protected]>

Author: julieqiu

internal/fetch/internal/doc/example.go

@@ -62,9 +62,6 @@ func Examples(testFiles ...*ast.File) []*Example {
 			if !ok || f.Recv != nil {
 				continue
 			}
-			if params := f.Type.Params; len(params.List) != 0 {
-				continue // function has params; not a valid example
-			}
 			numDecl++
 			name := f.Name.Name
 			if isTest(name, "Test") || isTest(name, "Benchmark") {
@@ -74,6 +71,9 @@ func Examples(testFiles ...*ast.File) []*Example {
 			if !isTest(name, "Example") {
 				continue
 			}
+			if params := f.Type.Params; len(params.List) != 0 {
+				continue // function has params; not a valid example
+			}
 			if f.Body == nil { // ast.File.Body nil dereference (see issue 28044)
 				continue
 			}

internal/fetch/internal/doc/example_test.go

@@ -331,26 +331,65 @@ func main() {
 	fmt.Println("Hello, world!")
 }
 `
+const exampleWholeFileFunction = `package foo_test
+
+func Foo(x int) {
+}
+
+func Example() {
+	fmt.Println("Hello, world!")
+	// Output: Hello, world!
+}
+`
+
+const exampleWholeFileFunctionOutput = `package main
+
+func Foo(x int) {
+}
+
+func main() {
+	fmt.Println("Hello, world!")
+}
+`
+
+var exampleWholeFileTestCases = []struct {
+	Title, Source, Play, Output string
+}{
+	{
+		"Methods",
+		exampleWholeFile,
+		exampleWholeFileOutput,
+		"Hello, world!\n",
+	},
+	{
+		"Function",
+		exampleWholeFileFunction,
+		exampleWholeFileFunctionOutput,
+		"Hello, world!\n",
+	},
+}
 
 func TestExamplesWholeFile(t *testing.T) {
-	fset := token.NewFileSet()
-	file, err := parser.ParseFile(fset, "test.go", strings.NewReader(exampleWholeFile), parser.ParseComments)
-	if err != nil {
-		t.Fatal(err)
-	}
-	es := doc.Examples(file)
-	if len(es) != 1 {
-		t.Fatalf("wrong number of examples; got %d want 1", len(es))
-	}
-	e := es[0]
-	if e.Name != "" {
-		t.Errorf("got Name == %q, want %q", e.Name, "")
-	}
-	if g, w := formatFile(t, fset, e.Play), exampleWholeFileOutput; g != w {
-		t.Errorf("got Play == %q, want %q", g, w)
-	}
-	if g, w := e.Output, "Hello, world!\n"; g != w {
-		t.Errorf("got Output == %q, want %q", g, w)
+	for _, c := range exampleWholeFileTestCases {
+		fset := token.NewFileSet()
+		file, err := parser.ParseFile(fset, "test.go", strings.NewReader(c.Source), parser.ParseComments)
+		if err != nil {
+			t.Fatal(err)
+		}
+		es := doc.Examples(file)
+		if len(es) != 1 {
+			t.Fatalf("%s: wrong number of examples; got %d want 1", c.Title, len(es))
+		}
+		e := es[0]
+		if e.Name != "" {
+			t.Errorf("%s: got Name == %q, want %q", c.Title, e.Name, "")
+		}
+		if g, w := formatFile(t, fset, e.Play), c.Play; g != w {
+			t.Errorf("%s: got Play == %q, want %q", c.Title, g, w)
+		}
+		if g, w := e.Output, c.Output; g != w {
+			t.Errorf("%s: got Output == %q, want %q", c.Title, g, w)
+		}
 	}
 }