1 files changed, 15 insertions, 16 deletions
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/ConfiguredAttributeMapperTest.java b/src/test/java/com/google/devtools/build/lib/analysis/ConfiguredAttributeMapperTest.java
index 15a90a87a1..f3014f20d0 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/ConfiguredAttributeMapperTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/ConfiguredAttributeMapperTest.java
@@ -19,13 +19,13 @@ import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.analysis.config.CompilationMode;
 import com.google.devtools.build.lib.analysis.util.BuildViewTestCase;
 import com.google.devtools.build.lib.cmdline.Label;
-import com.google.devtools.build.lib.packages.Attribute;
 import com.google.devtools.build.lib.packages.AttributeMap;
 import com.google.devtools.build.lib.packages.BuildType;
 import com.google.devtools.build.lib.packages.ConfiguredAttributeMapper;
 import com.google.devtools.build.lib.skyframe.ConfiguredTargetAndData;
 import com.google.devtools.build.lib.syntax.Type;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -114,37 +114,36 @@ public class ConfiguredAttributeMapperTest extends BuildViewTestCase {
         "    name = 'defaultdep',",
         "    srcs = ['defaultdep.sh'])");
 
-    final List<Label> visitedLabels = new ArrayList<>();
-    AttributeMap.AcceptsLabelAttribute testVisitor =
-        new AttributeMap.AcceptsLabelAttribute() {
-          @Override
-          public void acceptLabelAttribute(Label label, Attribute attribute) {
-            if (label.toString().contains("//a:")) { // Ignore implicit common dependencies.
-              visitedLabels.add(label);
-            }
-          }
-        };
-
-    final Label binSrc = Label.parseAbsolute("//a:bin.sh", ImmutableMap.of());
+    List<Label> visitedLabels = new ArrayList<>();
+    Label binSrc = Label.parseAbsolute("//a:bin.sh", ImmutableMap.of());
 
     useConfiguration("--define", "mode=a");
-    getMapper("//a:bin").visitLabels(testVisitor);
+    addRelevantLabels(getMapper("//a:bin").visitLabels(), visitedLabels);
     assertThat(visitedLabels)
         .containsExactly(binSrc, Label.parseAbsolute("//a:adep", ImmutableMap.of()));
 
     visitedLabels.clear();
     useConfiguration("--define", "mode=b");
-    getMapper("//a:bin").visitLabels(testVisitor);
+    addRelevantLabels(getMapper("//a:bin").visitLabels(), visitedLabels);
     assertThat(visitedLabels)
         .containsExactly(binSrc, Label.parseAbsolute("//a:bdep", ImmutableMap.of()));
 
     visitedLabels.clear();
     useConfiguration("--define", "mode=c");
-    getMapper("//a:bin").visitLabels(testVisitor);
+    addRelevantLabels(getMapper("//a:bin").visitLabels(), visitedLabels);
     assertThat(visitedLabels)
         .containsExactly(binSrc, Label.parseAbsolute("//a:defaultdep", ImmutableMap.of()));
   }
 
+  private static void addRelevantLabels(
+      Collection<AttributeMap.DepEdge> depEdges, Collection<Label> visitedLabels) {
+    depEdges
+        .stream()
+        .map(AttributeMap.DepEdge::getLabel)
+        .filter((label) -> label.getPackageIdentifier().getPackageFragment().toString().equals("a"))
+        .forEach(visitedLabels::add);
+  }
+
   /**
    * Tests that for configurable attributes where the *values* are evaluated in different
    * configurations, the configuration checking still uses the original configuration.