Add CONFIG to TestPathLookup

Author: prdoyle

test/framework/src/main/java/org/elasticsearch/bootstrap/BootstrapForTesting.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.common.network.IfConfig;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.Booleans;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.PathUtils;
 import org.elasticsearch.entitlement.bootstrap.TestEntitlementBootstrap;
 import org.elasticsearch.jdk.JarHell;
@@ -75,12 +76,20 @@ public class BootstrapForTesting {
 
         // Fire up entitlements
         try {
-            TestEntitlementBootstrap.bootstrap(javaTmpDir);
+            TestEntitlementBootstrap.bootstrap(javaTmpDir, maybePath(System.getProperty("tests.config")));
         } catch (IOException e) {
             throw new IllegalStateException(e.getClass().getSimpleName() + " while initializing entitlements for tests", e);
         }
     }
 
+    private static @Nullable Path maybePath(String str) {
+        if (str == null) {
+            return null;
+        } else {
+            return Path.of(str);
+        }
+    }
+
     // does nothing, just easy way to make sure the class is loaded.
     public static void ensureInitialized() {}
 }

test/framework/src/main/java/org/elasticsearch/entitlement/bootstrap/TestEntitlementBootstrap.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.bootstrap.TestBuildInfo;
 import org.elasticsearch.bootstrap.TestBuildInfoParser;
 import org.elasticsearch.bootstrap.TestScopeResolver;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Strings;
 import org.elasticsearch.core.SuppressForbidden;
 import org.elasticsearch.entitlement.initialization.EntitlementInitialization;
@@ -38,6 +39,9 @@
 import java.util.TreeSet;
 import java.util.stream.Collectors;
 
+import static org.elasticsearch.entitlement.runtime.policy.PathLookup.BaseDir.CONFIG;
+import static org.elasticsearch.entitlement.runtime.policy.PathLookup.BaseDir.TEMP;
+
 public class TestEntitlementBootstrap {
 
     private static final Logger logger = LogManager.getLogger(TestEntitlementBootstrap.class);
@@ -47,17 +51,25 @@ public class TestEntitlementBootstrap {
     /**
      * Activates entitlement checking in tests.
      */
-    public static void bootstrap(Path tempDir) throws IOException {
+    public static void bootstrap(@Nullable Path tempDir, @Nullable Path configDir) throws IOException {
         if (isEnabledForTest() == false) {
             return;
         }
-        TestPathLookup pathLookup = new TestPathLookup(List.of(tempDir));
+        TestPathLookup pathLookup = new TestPathLookup(Map.of(TEMP, zeroOrOne(tempDir), CONFIG, zeroOrOne(configDir)));
         policyManager = createPolicyManager(pathLookup);
         EntitlementInitialization.initializeArgs = new EntitlementInitialization.InitializeArgs(pathLookup, Set.of(), policyManager);
         logger.debug("Loading entitlement agent");
         EntitlementBootstrap.loadAgent(EntitlementBootstrap.findAgentJar(), EntitlementInitialization.class.getName());
     }
 
+    private static <T> List<T> zeroOrOne(T item) {
+        if (item == null) {
+            return List.of();
+        } else {
+            return List.of(item);
+        }
+    }
+
     public static boolean isEnabledForTest() {
         return Boolean.getBoolean("es.entitlement.enableForTests");
     }

test/framework/src/main/java/org/elasticsearch/entitlement/runtime/policy/TestPathLookup.java

@@ -10,13 +10,15 @@
 package org.elasticsearch.entitlement.runtime.policy;
 
 import java.nio.file.Path;
+import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.stream.Stream;
 
 public class TestPathLookup implements PathLookup {
-    final List<Path> tempDirPaths;
+    final Map<BaseDir, Collection<Path>> tempDirPaths;
 
-    public TestPathLookup(List<Path> tempDirPaths) {
+    public TestPathLookup(Map<BaseDir, Collection<Path>> tempDirPaths) {
         this.tempDirPaths = tempDirPaths;
     }
 
@@ -27,10 +29,7 @@ public Path pidFile() {
 
     @Override
     public Stream<Path> getBaseDirPaths(BaseDir baseDir) {
-        return switch (baseDir) {
-            case TEMP -> tempDirPaths.stream();
-            default -> Stream.empty();
-        };
+        return tempDirPaths.getOrDefault(baseDir, List.of()).stream();
     }
 
     @Override

test/framework/src/test/java/org/elasticsearch/entitlement/runtime/policy/TestPolicyManagerTests.java

@@ -31,7 +31,7 @@ public void setupPolicyManager() {
             Map.of(),
             c -> new PolicyScope(PLUGIN, "example-plugin" + scopeCounter.incrementAndGet(), "org.example.module"),
             Map.of(),
-            new TestPathLookup(List.of()),
+            new TestPathLookup(Map.of()),
             List.of()
         );
         policyManager.setActive(true);