Add test utility for getting all keys in the graph.

PiperOrigin-RevId: 623547996 Change-Id: Ib099b3ecf1521fa5815c471d7900cd86b22f8979
fmeum · Apr 10, 2024 · 85d5dfe · 85d5dfe
1 parent cf66672
commit 85d5dfe
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/src/test/java/com/google/devtools/build/lib/buildtool/GenQueryIntegrationTest.java b/src/test/java/com/google/devtools/build/lib/buildtool/GenQueryIntegrationTest.java
@@ -893,9 +893,7 @@ private void assertPartialQueryResult(String queryTarget, String... expected) th
   private String getQueryResult(String queryTarget) throws Exception {
     buildTarget(queryTarget);
     Artifact output = Iterables.getOnlyElement(getArtifacts(queryTarget));
-    assertThat(
-            getSkyframeExecutor().getEvaluator().getValues().keySet().stream()
-                .anyMatch(key -> key instanceof TransitiveTargetKey))
+    assertThat(getAllKeysInGraph().stream().anyMatch(key -> key instanceof TransitiveTargetKey))
         .isEqualTo(!ttvFree);
     return readContentAsLatin1String(output);
   }

diff --git a/src/test/java/com/google/devtools/build/lib/buildtool/util/BUILD b/src/test/java/com/google/devtools/build/lib/buildtool/util/BUILD
@@ -93,6 +93,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/com/google/devtools/build/lib/worker:worker_module",
         "//src/main/java/com/google/devtools/build/lib/worker:worker_process_metrics",
+        "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
         "//src/main/java/com/google/devtools/common/options",
         "//src/main/java/com/google/devtools/common/options:invocation_policy",
         "//src/main/protobuf:failure_details_java_proto",

diff --git a/src/test/java/com/google/devtools/build/lib/buildtool/util/BuildIntegrationTestCase.java b/src/test/java/com/google/devtools/build/lib/buildtool/util/BuildIntegrationTestCase.java
@@ -134,6 +134,7 @@
 import com.google.devtools.build.lib.vfs.util.FileSystems;
 import com.google.devtools.build.lib.worker.WorkerModule;
 import com.google.devtools.build.skyframe.NotifyingHelper;
+import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.common.options.OptionsBase;
 import com.google.devtools.common.options.OptionsParser;
 import com.google.devtools.common.options.OptionsParsingResult;
@@ -1232,4 +1233,8 @@ public synchronized void close() {
   protected Formatter getFormatterForLogging() {
     return new SimpleFormatter();
   }
+
+  protected Set<SkyKey> getAllKeysInGraph() {
+    return getSkyframeExecutor().getEvaluator().getValues().keySet();
+  }
 }