internal/fetchdatasource: handle build contexts properly

GetUnit returns a Unit with its Documentation set to a matching
BuildContext.  This is the same behavior as postgres.DB.GetUnit.

For golang/go#47780

Change-Id: I2bc23b7bc5a006e78bec54f6f3229e59ab5a03ef
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/345269
Trust: Jonathan Amsterdam <jba@google.com>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
TryBot-Result: kokoro <noreply+kokoro@google.com>
Reviewed-by: Jamal Carvalho <jamal@golang.org>
Reviewed-by: Julie Qiu <julie@golang.org>
diff --git a/internal/fetchdatasource/fetchdatasource.go b/internal/fetchdatasource/fetchdatasource.go
index fe2eacc..30115a4 100644
--- a/internal/fetchdatasource/fetchdatasource.go
+++ b/internal/fetchdatasource/fetchdatasource.go
@@ -199,10 +199,20 @@
 	if err != nil {
 		return nil, err
 	}
-	if u := findUnit(m, um.Path); u != nil {
-		return u, nil
+	u := findUnit(m, um.Path)
+	if u == nil {
+		return nil, fmt.Errorf("import path %s not found in module %s: %w", um.Path, um.ModulePath, derrors.NotFound)
 	}
-	return nil, fmt.Errorf("import path %s not found in module %s: %w", um.Path, um.ModulePath, derrors.NotFound)
+	// Return only the Documentation matching the given BuildContext, if any.
+	// Since we cache the module and its units, we have to copy this unit before we modify it.
+	// It can be a shallow copy, since we're only modifying the Unit.Documentation field.
+	u2 := *u
+	if d := matchingDoc(u.Documentation, bc); d != nil {
+		u2.Documentation = []*internal.Documentation{d}
+	} else {
+		u2.Documentation = nil
+	}
+	return &u2, nil
 }
 
 // findUnit returns the unit with the given path in m, or nil if none.
@@ -215,6 +225,23 @@
 	return nil
 }
 
+// matchingDoc returns the Documentation that matches the given build context
+// and comes earliest in build-context order. It returns nil if there is none.
+func matchingDoc(docs []*internal.Documentation, bc internal.BuildContext) *internal.Documentation {
+	var (
+		dMin  *internal.Documentation
+		bcMin = internal.BuildContext{GOOS: "unk", GOARCH: "unk"} // sorts last
+	)
+	for _, d := range docs {
+		dbc := d.BuildContext()
+		if bc.Match(dbc) && internal.CompareBuildContexts(dbc, bcMin) < 0 {
+			dMin = d
+			bcMin = dbc
+		}
+	}
+	return dMin
+}
+
 // GetLatestInfo returns latest information for unitPath and modulePath.
 func (ds *FetchDataSource) GetLatestInfo(ctx context.Context, unitPath, modulePath string, latestUnitMeta *internal.UnitMeta) (latest internal.LatestInfo, err error) {
 	defer derrors.Wrap(&err, "FetchDataSource.GetLatestInfo(ctx, %q, %q)", unitPath, modulePath)
diff --git a/internal/fetchdatasource/fetchdatasource_test.go b/internal/fetchdatasource/fetchdatasource_test.go
index 995dfad..e807ec7 100644
--- a/internal/fetchdatasource/fetchdatasource_test.go
+++ b/internal/fetchdatasource/fetchdatasource_test.go
@@ -449,7 +449,7 @@
 	}
 }
 
-func TestLocalGetUnit(t *testing.T) {
+func TestGetUnit(t *testing.T) {
 	// This is a simple test to verify that data is fetched correctly. The
 	// return value of FetchResult is tested in internal/fetch so no need
 	// to repeat it.
@@ -509,6 +509,45 @@
 	}
 }
 
+func TestBuildConstraints(t *testing.T) {
+	// The Unit returned by GetUnit should have a single Documentation that
+	// matches the BuildContext argument.
+	ctx, ds, teardown := setup(t, defaultTestModules, true)
+	defer teardown()
+
+	um := &internal.UnitMeta{
+		Path: "example.com/build-constraints/cpu",
+		ModuleInfo: internal.ModuleInfo{
+			ModulePath: "example.com/build-constraints",
+			Version:    version.Latest,
+		},
+	}
+	for _, test := range []struct {
+		in, want internal.BuildContext
+	}{
+		{internal.BuildContext{}, internal.BuildContextLinux},
+		{internal.BuildContextLinux, internal.BuildContextLinux},
+		{internal.BuildContextDarwin, internal.BuildContextDarwin},
+		{internal.BuildContext{GOOS: "LiverPaté", GOARCH: "DeTriomphe"}, internal.BuildContext{}},
+	} {
+		t.Run(test.in.String(), func(t *testing.T) {
+			u, err := ds.GetUnit(ctx, um, internal.AllFields, test.in)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if test.want == (internal.BuildContext{}) {
+				if len(u.Documentation) != 0 {
+					t.Error("got docs, want none")
+				}
+			} else if n := len(u.Documentation); n != 1 {
+				t.Errorf("got %d docs, want 1", n)
+			} else if got := u.Documentation[0].BuildContext(); got != test.want {
+				t.Errorf("got %s, want %s", got, test.want)
+			}
+		})
+	}
+}
+
 func TestCache(t *testing.T) {
 	ds := Options{}.New()
 	m1 := &internal.Module{}