1 files changed, 13 insertions, 20 deletions
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/select/AbstractAttributeMapperTest.java b/src/test/java/com/google/devtools/build/lib/analysis/select/AbstractAttributeMapperTest.java
index ac89f9f0eb..4901884590 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/select/AbstractAttributeMapperTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/select/AbstractAttributeMapperTest.java
@@ -16,7 +16,6 @@ package com.google.devtools.build.lib.analysis.select;
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.fail;
 
-import com.google.common.collect.Lists;
 import com.google.devtools.build.lib.analysis.util.BuildViewTestCase;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.packages.AbstractAttributeMapper;
@@ -29,6 +28,7 @@ import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.RuleClass;
 import com.google.devtools.build.lib.syntax.Type;
 import java.util.List;
+import java.util.stream.Collectors;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -116,27 +116,20 @@ public class AbstractAttributeMapperTest extends BuildViewTestCase {
     assertThat(mapper.isAttributeValueExplicitlySpecified("nonsense")).isFalse();
   }
 
-  protected static class VisitationRecorder implements AttributeMap.AcceptsLabelAttribute {
-    public List<String> labelsVisited = Lists.newArrayList();
-    private final String attrName;
-
-    public VisitationRecorder(String attrName) {
-      this.attrName = attrName;
-    }
-
-    @Override
-    public void acceptLabelAttribute(Label label, Attribute attribute) {
-      if (attribute.getName().equals(attrName)) {
-        labelsVisited.add(label.toString());
-      }
-    }
-  }
-
   @Test
   public void testVisitation() throws Exception {
-    VisitationRecorder recorder = new VisitationRecorder("srcs");
-    mapper.visitLabels(recorder);
-    assertThat(recorder.labelsVisited).containsExactly("//p:a", "//p:b", "//p:c");
+    assertThat(getLabelsForAttribute(mapper, "srcs")).containsExactly("//p:a", "//p:b", "//p:c");
+  }
+
+  protected static List<String> getLabelsForAttribute(
+      AttributeMap attributeMap, String attributeName) throws InterruptedException {
+    return attributeMap
+        .visitLabels()
+        .stream()
+        .filter((d) -> d.getAttribute().getName().equals(attributeName))
+        .map(AttributeMap.DepEdge::getLabel)
+        .map(Label::toString)
+        .collect(Collectors.toList());
   }
 
   @Test