From 85d5dfeaff41a77de6f53c03215ef1faa32d5a15 Mon Sep 17 00:00:00 2001
From: Googler <noreply@google.com>
Date: Wed, 10 Apr 2024 11:04:35 -0700
Subject: [PATCH] Add test utility for getting all keys in the graph.

PiperOrigin-RevId: 623547996
Change-Id: Ib099b3ecf1521fa5815c471d7900cd86b22f8979
---
 .../build/lib/buildtool/GenQueryIntegrationTest.java         | 4 +---
 .../java/com/google/devtools/build/lib/buildtool/util/BUILD  | 1 +
 .../build/lib/buildtool/util/BuildIntegrationTestCase.java   | 5 +++++
 3 files changed, 7 insertions(+), 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
index 4c0afb2dcc7a8e..429bdb3e0e7293 100644
--- 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
index 22327bf0dbbe0d..cbb026dab9072c 100644
--- 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
index f894601c6e26b3..b696ca6085c5a3 100644
--- 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();
+  }
 }