Add TestServices.newCallContext

runtime/service/src/test/java/org/apache/polaris/service/admin/ManagementServiceTest.java

@@ -170,15 +170,6 @@ private PolarisMetaStoreManager setupMetaStoreManager() {
     return metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext);
   }
 
-  private PolarisCallContext setupCallContext() {
-    MetaStoreManagerFactory metaStoreManagerFactory = services.metaStoreManagerFactory();
-    RealmContext realmContext = services.realmContext();
-    return new PolarisCallContext(
-        realmContext,
-        metaStoreManagerFactory.getOrCreateSession(realmContext),
-        services.polarisDiagnostics());
-  }
-
   private PolarisAdminService setupPolarisAdminService(
       PolarisMetaStoreManager metaStoreManager, PolarisCallContext callContext) {
     return new PolarisAdminService(
@@ -252,7 +243,7 @@ private PrincipalRoleEntity createRole(
   @Test
   public void testCannotAssignFederatedEntities() {
     PolarisMetaStoreManager metaStoreManager = setupMetaStoreManager();
-    PolarisCallContext callContext = setupCallContext();
+    PolarisCallContext callContext = services.newCallContext();
     PolarisAdminService polarisAdminService =
         setupPolarisAdminService(metaStoreManager, callContext);
 
@@ -272,7 +263,7 @@ public void testCannotAssignFederatedEntities() {
   @Test
   public void testCatalogNotReturnedWhenDeletedAfterListBeforeGet() {
     PolarisMetaStoreManager metaStoreManager = Mockito.spy(setupMetaStoreManager());
-    PolarisCallContext callContext = setupCallContext();
+    PolarisCallContext callContext = services.newCallContext();
     PolarisAdminService polarisAdminService =
         setupPolarisAdminService(metaStoreManager, callContext);
 

runtime/service/src/test/java/org/apache/polaris/service/catalog/io/FileIOFactoryTest.java

@@ -25,7 +25,6 @@
 import com.google.common.collect.ImmutableMap;
 import jakarta.annotation.Nonnull;
 import java.lang.reflect.Method;
-import java.time.Clock;
 import java.util.List;
 import java.util.Map;
 import org.apache.iceberg.Schema;
@@ -35,7 +34,6 @@
 import org.apache.iceberg.inmemory.InMemoryFileIO;
 import org.apache.iceberg.io.FileIO;
 import org.apache.iceberg.types.Types;
-import org.apache.polaris.core.PolarisCallContext;
 import org.apache.polaris.core.admin.model.AwsStorageConfigInfo;
 import org.apache.polaris.core.admin.model.Catalog;
 import org.apache.polaris.core.admin.model.CatalogProperties;
@@ -136,13 +134,7 @@ FileIO loadFileIOInternal(
             .fileIOFactorySupplier(fileIOFactorySupplier)
             .build();
 
-    callContext =
-        new PolarisCallContext(
-            realmContext,
-            testServices.metaStoreManagerFactory().getOrCreateSession(realmContext),
-            testServices.polarisDiagnostics(),
-            testServices.configurationStore(),
-            Clock.systemUTC());
+    callContext = testServices.newCallContext();
   }
 
   @AfterEach

runtime/service/src/test/java/org/apache/polaris/service/task/BatchFileCleanupTaskHandlerTest.java

@@ -43,30 +43,33 @@
 import org.apache.iceberg.inmemory.InMemoryFileIO;
 import org.apache.iceberg.io.FileIO;
 import org.apache.polaris.core.PolarisCallContext;
-import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
+import org.apache.polaris.core.PolarisDiagnostics;
 import org.apache.polaris.core.context.RealmContext;
 import org.apache.polaris.core.entity.AsyncTaskType;
 import org.apache.polaris.core.entity.TaskEntity;
+import org.apache.polaris.core.persistence.BasePersistence;
 import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
 import org.apache.polaris.service.TestFileIOFactory;
 import org.junit.jupiter.api.Test;
 
 @QuarkusTest
 public class BatchFileCleanupTaskHandlerTest {
+  @Inject PolarisDiagnostics diagnostics;
   @Inject MetaStoreManagerFactory metaStoreManagerFactory;
   private final RealmContext realmContext = () -> "realmName";
 
   private TaskFileIOSupplier buildTaskFileIOSupplier(FileIO fileIO) {
     return new TaskFileIOSupplier(new TestFileIOFactory(fileIO));
   }
 
+  private PolarisCallContext newCallContext() {
+    BasePersistence metaStore = metaStoreManagerFactory.getOrCreateSession(realmContext);
+    return new PolarisCallContext(realmContext, metaStore, diagnostics);
+  }
+
   @Test
   public void testMetadataFileCleanup() throws IOException {
-    PolarisCallContext polarisCallContext =
-        new PolarisCallContext(
-            realmContext,
-            metaStoreManagerFactory.getOrCreateSession(realmContext),
-            new PolarisDefaultDiagServiceImpl());
+    PolarisCallContext polarisCallContext = newCallContext();
     FileIO fileIO =
         new InMemoryFileIO() {
           @Override
@@ -175,11 +178,7 @@ public void close() {
 
   @Test
   public void testMetadataFileCleanupIfFileNotExist() throws IOException {
-    PolarisCallContext polarisCallContext =
-        new PolarisCallContext(
-            realmContext,
-            metaStoreManagerFactory.getOrCreateSession(realmContext),
-            new PolarisDefaultDiagServiceImpl());
+    PolarisCallContext polarisCallContext = newCallContext();
     FileIO fileIO = new InMemoryFileIO();
     TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1");
     BatchFileCleanupTaskHandler handler =
@@ -219,11 +218,7 @@ public void testMetadataFileCleanupIfFileNotExist() throws IOException {
 
   @Test
   public void testCleanupWithRetries() throws IOException {
-    PolarisCallContext polarisCallContext =
-        new PolarisCallContext(
-            realmContext,
-            metaStoreManagerFactory.getOrCreateSession(realmContext),
-            new PolarisDefaultDiagServiceImpl());
+    PolarisCallContext polarisCallContext = newCallContext();
     Map<String, AtomicInteger> retryCounter = new HashMap<>();
     FileIO fileIO =
         new InMemoryFileIO() {

runtime/service/src/test/java/org/apache/polaris/service/task/ManifestFileCleanupTaskHandlerTest.java

@@ -41,16 +41,18 @@
 import org.apache.iceberg.io.OutputFile;
 import org.apache.iceberg.io.PositionOutputStream;
 import org.apache.polaris.core.PolarisCallContext;
-import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
+import org.apache.polaris.core.PolarisDiagnostics;
 import org.apache.polaris.core.context.RealmContext;
 import org.apache.polaris.core.entity.AsyncTaskType;
 import org.apache.polaris.core.entity.TaskEntity;
+import org.apache.polaris.core.persistence.BasePersistence;
 import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
 import org.apache.polaris.service.TestFileIOFactory;
 import org.junit.jupiter.api.Test;
 
 @QuarkusTest
 class ManifestFileCleanupTaskHandlerTest {
+  @Inject PolarisDiagnostics diagnostics;
   @Inject MetaStoreManagerFactory metaStoreManagerFactory;
 
   private final RealmContext realmContext = () -> "realmName";
@@ -59,13 +61,14 @@ private TaskFileIOSupplier buildTaskFileIOSupplier(FileIO fileIO) {
     return new TaskFileIOSupplier(new TestFileIOFactory(fileIO));
   }
 
+  private PolarisCallContext newCallContext() {
+    BasePersistence metaStore = metaStoreManagerFactory.getOrCreateSession(realmContext);
+    return new PolarisCallContext(realmContext, metaStore, diagnostics);
+  }
+
   @Test
   public void testCleanupFileNotExists() throws IOException {
-    PolarisCallContext polarisCallContext =
-        new PolarisCallContext(
-            realmContext,
-            metaStoreManagerFactory.getOrCreateSession(realmContext),
-            new PolarisDefaultDiagServiceImpl());
+    PolarisCallContext polarisCallContext = newCallContext();
     FileIO fileIO = new InMemoryFileIO();
     TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1");
 
@@ -91,11 +94,7 @@ public void testCleanupFileNotExists() throws IOException {
 
   @Test
   public void testCleanupFileManifestExistsDataFilesDontExist() throws IOException {
-    PolarisCallContext polarisCallContext =
-        new PolarisCallContext(
-            realmContext,
-            metaStoreManagerFactory.getOrCreateSession(realmContext),
-            new PolarisDefaultDiagServiceImpl());
+    PolarisCallContext polarisCallContext = newCallContext();
     FileIO fileIO = new InMemoryFileIO();
     TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1");
     ManifestFileCleanupTaskHandler handler =
@@ -119,11 +118,7 @@ public void testCleanupFileManifestExistsDataFilesDontExist() throws IOException
 
   @Test
   public void testCleanupFiles() throws IOException {
-    PolarisCallContext polarisCallContext =
-        new PolarisCallContext(
-            realmContext,
-            metaStoreManagerFactory.getOrCreateSession(realmContext),
-            new PolarisDefaultDiagServiceImpl());
+    PolarisCallContext polarisCallContext = newCallContext();
     FileIO fileIO =
         new InMemoryFileIO() {
           @Override
@@ -164,11 +159,7 @@ public void close() {
 
   @Test
   public void testCleanupFilesWithRetries() throws IOException {
-    PolarisCallContext polarisCallContext =
-        new PolarisCallContext(
-            realmContext,
-            metaStoreManagerFactory.getOrCreateSession(realmContext),
-            new PolarisDefaultDiagServiceImpl());
+    PolarisCallContext polarisCallContext = newCallContext();
     Map<String, AtomicInteger> retryCounter = new HashMap<>();
     FileIO fileIO =
         new InMemoryFileIO() {

runtime/service/src/test/java/org/apache/polaris/service/task/TaskExecutorImplTest.java

@@ -22,7 +22,6 @@
 import org.apache.polaris.core.context.CallContext;
 import org.apache.polaris.core.context.RealmContext;
 import org.apache.polaris.core.entity.TaskEntity;
-import org.apache.polaris.core.persistence.BasePersistence;
 import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
 import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
 import org.apache.polaris.service.TestServices;
@@ -47,10 +46,8 @@ void testEventsAreEmitted() {
     MetaStoreManagerFactory metaStoreManagerFactory = testServices.metaStoreManagerFactory();
     PolarisMetaStoreManager metaStoreManager =
         metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext);
-    BasePersistence bp = metaStoreManagerFactory.getOrCreateSession(realmContext);
 
-    PolarisCallContext polarisCallCtx =
-        new PolarisCallContext(realmContext, bp, testServices.polarisDiagnostics());
+    PolarisCallContext polarisCallCtx = testServices.newCallContext();
 
     // This task doesn't have a type so it won't be handle-able by a real handler. We register a
     // test TaskHandler below that can handle any task.

runtime/service/src/testFixtures/java/org/apache/polaris/service/TestServices.java

@@ -297,4 +297,10 @@ public String getAuthenticationScheme() {
           polarisEventListener);
     }
   }
+
+  public PolarisCallContext newCallContext() {
+    BasePersistence metaStore = metaStoreManagerFactory.getOrCreateSession(realmContext);
+    return new PolarisCallContext(
+        realmContext, metaStore, polarisDiagnostics, configurationStore, Clock.systemUTC());
+  }
 }