From 9feecad23b81489aab72873ed43f383a2abfc9ac Mon Sep 17 00:00:00 2001
From: Dom Narducci <dom@uber.com>
Date: Wed, 5 Oct 2022 16:32:48 -0700
Subject: [PATCH] language/go should consider default_visibility set by
 OtherGen (#783)

Previously, the implementation of the language/go extension's GenerateRules method only considered args.File, and therefore missed scenarios where another extension (as recommended in the Issue) creates a package(default_visibility = ...)
---
 language/go/generate.go      | 17 ++++++++++++++++-
 language/go/generate_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/language/go/generate.go b/language/go/generate.go
index d817f6eb0..0b8cd9523 100644
--- a/language/go/generate.go
+++ b/language/go/generate.go
@@ -181,7 +181,7 @@ func (gl *goLang) GenerateRules(args language.GenerateArgs) language.GenerateRes
 	g := &generator{
 		c:                   c,
 		rel:                 args.Rel,
-		shouldSetVisibility: args.File == nil || !args.File.HasDefaultVisibility(),
+		shouldSetVisibility: shouldSetVisibility(args),
 	}
 	var res language.GenerateResult
 	var rules []*rule.Rule
@@ -776,3 +776,18 @@ func escapeOption(opt string) string {
 		"$", "$$",
 	).Replace(opt)
 }
+
+func shouldSetVisibility(args language.GenerateArgs) bool {
+	if args.File != nil && args.File.HasDefaultVisibility() {
+		return false
+	}
+
+	for _, r := range args.OtherGen {
+		// This is kind of the same test as *File.HasDefaultVisibility(),
+		// but for previously defined rules.
+		if r.Kind() == "package" && r.Attr("default_visibility") != nil {
+			return false
+		}
+	}
+	return true
+}
diff --git a/language/go/generate_test.go b/language/go/generate_test.go
index 4d087168c..95af7aa5e 100644
--- a/language/go/generate_test.go
+++ b/language/go/generate_test.go
@@ -262,6 +262,35 @@ func TestConsumedGenFiles(t *testing.T) {
 	}
 }
 
+// Test visibility attribute is only set if no default visibility is provided
+// by the file or other rules.
+func TestShouldSetVisibility(t *testing.T) {
+	if !shouldSetVisibility(language.GenerateArgs{}) {
+		t.Error("got 'False' for shouldSetVisibility with default args; expected 'True'")
+	}
+
+	if !shouldSetVisibility(language.GenerateArgs{
+		File: rule.EmptyFile("path", "pkg"),
+	}) {
+		t.Error("got 'False' for shouldSetVisibility with empty file; expected 'True'")
+	}
+
+	fileWithDefaultVisibile, _ := rule.LoadData("path", "pkg", []byte(`package(default_visibility = "//src:__subpackages__")`))
+	if shouldSetVisibility(language.GenerateArgs{
+		File: fileWithDefaultVisibile,
+	}) {
+		t.Error("got 'True' for shouldSetVisibility with file with default visibility; expected 'False'")
+	}
+
+	defaultVisibilityRule := rule.NewRule("package", "")
+	defaultVisibilityRule.SetAttr("default_visibility", []string{"//src:__subpackages__"})
+	if shouldSetVisibility(language.GenerateArgs{
+		OtherGen: []*rule.Rule{defaultVisibilityRule},
+	}) {
+		t.Error("got 'True' for shouldSetVisibility with rule defining a default visibility; expected 'False'")
+	}
+}
+
 func prebuiltProtoRules() []*rule.Rule {
 	protoRule := rule.NewRule("proto_library", "foo_proto")
 	protoRule.SetAttr("srcs", []string{"foo.proto"})