From 598f5b93778b2e2e75265ae54484dd54477433f5 Mon Sep 17 00:00:00 2001
From: Tyler Bui-Palsulich <26876514+tbpg@users.noreply.github.com>
Date: Tue, 22 Jun 2021 13:55:43 -0400
Subject: [PATCH] feat(internal/godocfx): different metadata for different
 modules (#4297)

---
 internal/godocfx/go.mod          |  1 +
 internal/godocfx/godocfx_test.go | 56 ++++++++++++++++++++++++++++++++
 internal/godocfx/main.go         | 26 ++++++++++++---
 3 files changed, 78 insertions(+), 5 deletions(-)

diff --git a/internal/godocfx/go.mod b/internal/godocfx/go.mod
index a958cb2c4ed..4a36b56fe3c 100644
--- a/internal/godocfx/go.mod
+++ b/internal/godocfx/go.mod
@@ -7,6 +7,7 @@ require (
 	cloud.google.com/go/bigquery v1.8.0
 	cloud.google.com/go/datastore v1.1.0
 	cloud.google.com/go/storage v1.11.0
+	github.com/google/go-cmp v0.5.6
 	github.com/yuin/goldmark v1.3.8
 	golang.org/x/tools v0.1.3
 	gopkg.in/yaml.v2 v2.4.0
diff --git a/internal/godocfx/godocfx_test.go b/internal/godocfx/godocfx_test.go
index d227b0b89bf..4b92a84a4d2 100644
--- a/internal/godocfx/godocfx_test.go
+++ b/internal/godocfx/godocfx_test.go
@@ -17,14 +17,19 @@
 package main
 
 import (
+	"bytes"
 	"flag"
+	"fmt"
 	"io/ioutil"
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 
 	_ "cloud.google.com/go/bigquery" // Implicitly required by test.
 	_ "cloud.google.com/go/storage"  // Implicitly required by test.
+	"github.com/google/go-cmp/cmp"
+	"golang.org/x/tools/go/packages"
 )
 
 var updateGoldens bool
@@ -198,3 +203,54 @@ func TestHasPrefix(t *testing.T) {
 		}
 	}
 }
+
+func TestWriteMetadata(t *testing.T) {
+	now := time.Now()
+
+	want := fmt.Sprintf(`update_time {
+	seconds: %d
+	nanos: %d
+}
+name: "cloud.google.com/go"
+version: "100.0.0"
+language: "go"
+`, now.Unix(), now.Nanosecond())
+
+	wantAppEngine := fmt.Sprintf(`update_time {
+	seconds: %d
+	nanos: %d
+}
+name: "google.golang.org/appengine/v2"
+version: "2.0.0"
+language: "go"
+stem: "/appengine/docs/standard/go/reference/services/bundled"
+`, now.Unix(), now.Nanosecond())
+
+	tests := []struct {
+		path    string
+		version string
+		want    string
+	}{
+		{
+			path:    "cloud.google.com/go",
+			version: "100.0.0",
+			want:    want,
+		},
+		{
+			path:    "google.golang.org/appengine/v2",
+			version: "2.0.0",
+			want:    wantAppEngine,
+		},
+	}
+	for _, test := range tests {
+		var buf bytes.Buffer
+		module := &packages.Module{
+			Path:    test.path,
+			Version: test.version,
+		}
+		writeMetadata(&buf, now, module)
+		if diff := cmp.Diff(test.want, buf.String()); diff != "" {
+			t.Errorf("writeMetadata(%q) got unexpected diff (-want +got):\n\n%s", test.path, diff)
+		}
+	}
+}
diff --git a/internal/godocfx/main.go b/internal/godocfx/main.go
index 73cf4092e56..44ad0a21e40 100644
--- a/internal/godocfx/main.go
+++ b/internal/godocfx/main.go
@@ -49,6 +49,7 @@ import (
 	"strings"
 	"time"
 
+	"golang.org/x/tools/go/packages"
 	"gopkg.in/yaml.v2"
 )
 
@@ -245,12 +246,27 @@ func write(outDir string, r *result) error {
 	}
 	defer f.Close()
 	now := time.Now().UTC()
-	fmt.Fprintf(f, `update_time {
-  seconds: %d
-  nanos: %d
+	writeMetadata(f, now, r.module)
+	return nil
+}
+
+func writeMetadata(w io.Writer, now time.Time, module *packages.Module) {
+	fmt.Fprintf(w, `update_time {
+	seconds: %d
+	nanos: %d
 }
 name: %q
 version: %q
-language: "go"`, now.Unix(), now.Nanosecond(), r.module.Path, r.module.Version)
-	return nil
+language: "go"
+`, now.Unix(), now.Nanosecond(), module.Path, module.Version)
+
+	// Some modules specify a different path to serve from.
+	// The URL will be /[stem]/[version]/[pkg path relative to module].
+	// Alternatively, we could plumb this through command line flags.
+	switch module.Path {
+	case "google.golang.org/appengine":
+		fmt.Fprintf(w, "stem: \"/appengine/docs/standard/go111/reference\"\n")
+	case "google.golang.org/appengine/v2":
+		fmt.Fprintf(w, "stem: \"/appengine/docs/standard/go/reference/services/bundled\"\n")
+	}
 }