review comments

Author: mosche

server/src/test/java/org/elasticsearch/bootstrap/EntitlementMetaTests.java

@@ -41,7 +41,7 @@
  */
 public class EntitlementMetaTests extends ESTestCase {
     public void testSelfTestPasses() {
-        assumeTrue("Not yet working in serverless", TestEntitlementsRule.isEnabledForTest());
+        assumeTrue("Not yet working in serverless", TestEntitlementsRule.isEnabledForTests());
         Elasticsearch.entitlementSelfTest();
     }
 

server/src/test/java/org/elasticsearch/bootstrap/WithEntitlementsOnTestCodeMetaTests.java

@@ -30,13 +30,13 @@ public class WithEntitlementsOnTestCodeMetaTests extends ESTestCase {
      * is called from server code. The self-test should pass as usual.
      */
     public void testSelfTestPasses() {
-        assumeTrue("Not yet working in serverless", TestEntitlementsRule.isEnabledForTest());
+        assumeTrue("Not yet working in serverless", TestEntitlementsRule.isEnabledForTests());
         Elasticsearch.entitlementSelfTest();
     }
 
     @SuppressForbidden(reason = "Testing that a forbidden API is disallowed")
     public void testForbiddenActionDenied() {
-        assumeTrue("Not yet working in serverless", TestEntitlementsRule.isEnabledForTest());
+        assumeTrue("Not yet working in serverless", TestEntitlementsRule.isEnabledForTests());
         assertThrows(NotEntitledException.class, () -> Path.of(".").toRealPath());
     }
 }

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

@@ -74,7 +74,7 @@ public class TestEntitlementsRule implements TestRule {
     static {
         PathLookup pathLookup = new TestPathLookup(BASE_DIR_PATHS);
         try {
-            if (isEnabledForTest()) {
+            if (isEnabledForTests()) {
                 POLICY_MANAGER = createPolicyManager(pathLookup);
                 loadAgent(POLICY_MANAGER, pathLookup);
             } else {
@@ -107,6 +107,7 @@ public Statement apply(Statement base, Description description) {
                 public void evaluate() throws Throwable {
                     if (active.compareAndSet(false, true)) {
                         try {
+                            BASE_DIR_PATHS.keySet().retainAll(List.of(TEMP));
                             POLICY_MANAGER.setActive(false == withoutEntitlements);
                             POLICY_MANAGER.setTriviallyAllowingTestCode(false == withEntitlementsOnTestCode);
                             if (entitledPackages != null) {
@@ -115,16 +116,12 @@ public void evaluate() throws Throwable {
                             } else {
                                 POLICY_MANAGER.setEntitledTestPackages();
                             }
-                            BASE_DIR_PATHS.keySet().retainAll(List.of(TEMP));
                             POLICY_MANAGER.clearModuleEntitlementsCache();
                             // evaluate the suite
                             base.evaluate();
                         } finally {
-                            POLICY_MANAGER.setActive(false);
-                            POLICY_MANAGER.setTriviallyAllowingTestCode(true);
-                            POLICY_MANAGER.setEntitledTestPackages();
                             BASE_DIR_PATHS.keySet().retainAll(List.of(TEMP));
-                            POLICY_MANAGER.clearModuleEntitlementsCache();
+                            POLICY_MANAGER.resetAfterTest();
                             active.set(false);
                         }
                     } else {
@@ -144,6 +141,7 @@ public void evaluate() throws Throwable {
     /**
      * Temporarily adds node paths based entitlements based on a node's {@code settings} and {@code configPath}
      * until the returned handle is closed.
+     * @see PathLookup
      */
     public Closeable addEntitledNodePaths(Settings settings, Path configPath) {
         if (POLICY_MANAGER == null) {
@@ -250,7 +248,7 @@ private static BiFunction<BaseDir, Collection<Path>, Collection<Path>> baseDirMo
         };
     }
 
-    public static boolean isEnabledForTest() {
+    public static boolean isEnabledForTests() {
         return Booleans.parseBoolean(System.getProperty("es.entitlement.enableForTests", "false"));
     }
 
@@ -272,7 +270,7 @@ private static TestPolicyManager createPolicyManager(PathLookup pathLookup) thro
 
         String separator = System.getProperty("path.separator");
 
-        // In productions, plugins would have access to their respective bundle directories,
+        // In production, plugins would have access to their respective bundle directories,
         // and so they'd be able to read from their jars. In testing, we approximate this
         // by considering the entire classpath to be "source paths" of all plugins. This
         // also has the effect of granting read access to everything on the test-only classpath,

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

@@ -55,6 +55,13 @@ public TestPolicyManager(
         this.testOnlyClasspath = testOnlyClasspath;
     }
 
+    public final void resetAfterTest() {
+        isActive = false;
+        isTriviallyAllowingTestCode = true;
+        entitledTestPackages = TEST_FRAMEWORK_PACKAGE_PREFIXES;
+        clearModuleEntitlementsCache();
+    }
+
     public void setActive(boolean newValue) {
         this.isActive = newValue;
     }