From 3364dce8cbe255abbaece149d0633158777df603 Mon Sep 17 00:00:00 2001
From: Jay Conrod <jay@engflow.com>
Date: Tue, 2 Jan 2024 17:38:53 -0800
Subject: [PATCH] make test work with and without bzlmod

---
 .../gopackagesdriver/gopackagesdriver_test.go | 28 +++++++++++--------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/go/tools/gopackagesdriver/gopackagesdriver_test.go b/go/tools/gopackagesdriver/gopackagesdriver_test.go
index e7aaba159..8fb71be3e 100644
--- a/go/tools/gopackagesdriver/gopackagesdriver_test.go
+++ b/go/tools/gopackagesdriver/gopackagesdriver_test.go
@@ -191,19 +191,25 @@ func TestIncompatible(t *testing.T) {
 		t.Fatalf("unmarshaling response: %v", err)
 	}
 
-	helloLabel := "@@//:hello"
-	incompatibleLabel := "@@//:incompatible"
-	var foundHello bool
+	rootLabels := make(map[string]bool)
 	for _, root := range resp.Roots {
-		if root == helloLabel {
-			foundHello = true
-		}
-		if root == incompatibleLabel {
-			t.Errorf("response contains root %s", incompatibleLabel)
-		}
+		rootLabels[root] = true
 	}
-	if !foundHello {
-		t.Errorf("response does not contain root %s; roots were %s", helloLabel, strings.Join(resp.Roots, ", "))
+
+	// Verify //:hello is in .Roots and check whether its label starts with
+	// "@@" (bzlmod) or "@" (not bzlmod).
+	var incompatibleLabel string
+	if rootLabels["@@//:hello"] {
+		incompatibleLabel = "@@//:incompatible"
+	} else if rootLabels["@//:hello"] {
+		incompatibleLabel = "@//:incompatible"
+	} else {
+		t.Fatalf("response does not contain //:hello; roots were %s", strings.Join(resp.Roots, ", "))
+	}
+
+	// Verify //:incompatible is NOT in .Roots.
+	if rootLabels[incompatibleLabel] {
+		t.Fatalf("response contains root %s", incompatibleLabel)
 	}
 }