1 files changed, 2 insertions, 6 deletions
diff --git a/src/test/java/com/google/devtools/build/lib/analysis/select/AggregatingAttributeMapperTest.java b/src/test/java/com/google/devtools/build/lib/analysis/select/AggregatingAttributeMapperTest.java
index a8f912602c..073c751766 100644
--- a/src/test/java/com/google/devtools/build/lib/analysis/select/AggregatingAttributeMapperTest.java
+++ b/src/test/java/com/google/devtools/build/lib/analysis/select/AggregatingAttributeMapperTest.java
@@ -146,9 +146,7 @@ public class AggregatingAttributeMapperTest extends AbstractAttributeMapperTest
         "              '" + BuildType.Selector.DEFAULT_CONDITION_KEY + "': ['default.sh'],",
         "          }))");
 
-    VisitationRecorder recorder = new VisitationRecorder("srcs");
-    AggregatingAttributeMapper.of(rule).visitLabels(recorder);
-    assertThat(recorder.labelsVisited)
+    assertThat(getLabelsForAttribute(AggregatingAttributeMapper.of(rule), "srcs"))
         .containsExactlyElementsIn(
             ImmutableList.of(
                 "//a:a.sh", "//a:b.sh", "//a:default.sh", "//conditions:a", "//conditions:b"));
@@ -163,9 +161,7 @@ public class AggregatingAttributeMapperTest extends AbstractAttributeMapperTest
         "        '//conditions:a': None,",
         "    }))");
 
-    VisitationRecorder recorder = new VisitationRecorder("malloc");
-    AggregatingAttributeMapper.of(rule).visitLabels(recorder);
-    assertThat(recorder.labelsVisited)
+    assertThat(getLabelsForAttribute(AggregatingAttributeMapper.of(rule), "malloc"))
         .containsExactly("//conditions:a", getDefaultMallocLabel(rule).toString());
   }