Skip to content

Commit f34adca

Browse files
XN137XN137
authored
Add TestServices.newCallContext (#2249)
also add local `newCallContext` helper in some test classes
1 parent 4e82cd1 commit f34adca

File tree

6 files changed

+33
-61
lines changed

6 files changed

+33
-61
lines changed

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

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -170,15 +170,6 @@ private PolarisMetaStoreManager setupMetaStoreManager() {
170170
return metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext);
171171
}
172172

173-
private PolarisCallContext setupCallContext() {
174-
MetaStoreManagerFactory metaStoreManagerFactory = services.metaStoreManagerFactory();
175-
RealmContext realmContext = services.realmContext();
176-
return new PolarisCallContext(
177-
realmContext,
178-
metaStoreManagerFactory.getOrCreateSession(realmContext),
179-
services.polarisDiagnostics());
180-
}
181-
182173
private PolarisAdminService setupPolarisAdminService(
183174
PolarisMetaStoreManager metaStoreManager, PolarisCallContext callContext) {
184175
return new PolarisAdminService(
@@ -252,7 +243,7 @@ private PrincipalRoleEntity createRole(
252243
@Test
253244
public void testCannotAssignFederatedEntities() {
254245
PolarisMetaStoreManager metaStoreManager = setupMetaStoreManager();
255-
PolarisCallContext callContext = setupCallContext();
246+
PolarisCallContext callContext = services.newCallContext();
256247
PolarisAdminService polarisAdminService =
257248
setupPolarisAdminService(metaStoreManager, callContext);
258249

@@ -272,7 +263,7 @@ public void testCannotAssignFederatedEntities() {
272263
@Test
273264
public void testCatalogNotReturnedWhenDeletedAfterListBeforeGet() {
274265
PolarisMetaStoreManager metaStoreManager = Mockito.spy(setupMetaStoreManager());
275-
PolarisCallContext callContext = setupCallContext();
266+
PolarisCallContext callContext = services.newCallContext();
276267
PolarisAdminService polarisAdminService =
277268
setupPolarisAdminService(metaStoreManager, callContext);
278269

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

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import com.google.common.collect.ImmutableMap;
2626
import jakarta.annotation.Nonnull;
2727
import java.lang.reflect.Method;
28-
import java.time.Clock;
2928
import java.util.List;
3029
import java.util.Map;
3130
import org.apache.iceberg.Schema;
@@ -35,7 +34,6 @@
3534
import org.apache.iceberg.inmemory.InMemoryFileIO;
3635
import org.apache.iceberg.io.FileIO;
3736
import org.apache.iceberg.types.Types;
38-
import org.apache.polaris.core.PolarisCallContext;
3937
import org.apache.polaris.core.admin.model.AwsStorageConfigInfo;
4038
import org.apache.polaris.core.admin.model.Catalog;
4139
import org.apache.polaris.core.admin.model.CatalogProperties;
@@ -136,13 +134,7 @@ FileIO loadFileIOInternal(
136134
.fileIOFactorySupplier(fileIOFactorySupplier)
137135
.build();
138136

139-
callContext =
140-
new PolarisCallContext(
141-
realmContext,
142-
testServices.metaStoreManagerFactory().getOrCreateSession(realmContext),
143-
testServices.polarisDiagnostics(),
144-
testServices.configurationStore(),
145-
Clock.systemUTC());
137+
callContext = testServices.newCallContext();
146138
}
147139

148140
@AfterEach

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

Lines changed: 11 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -43,30 +43,33 @@
4343
import org.apache.iceberg.inmemory.InMemoryFileIO;
4444
import org.apache.iceberg.io.FileIO;
4545
import org.apache.polaris.core.PolarisCallContext;
46-
import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
46+
import org.apache.polaris.core.PolarisDiagnostics;
4747
import org.apache.polaris.core.context.RealmContext;
4848
import org.apache.polaris.core.entity.AsyncTaskType;
4949
import org.apache.polaris.core.entity.TaskEntity;
50+
import org.apache.polaris.core.persistence.BasePersistence;
5051
import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
5152
import org.apache.polaris.service.TestFileIOFactory;
5253
import org.junit.jupiter.api.Test;
5354

5455
@QuarkusTest
5556
public class BatchFileCleanupTaskHandlerTest {
57+
@Inject PolarisDiagnostics diagnostics;
5658
@Inject MetaStoreManagerFactory metaStoreManagerFactory;
5759
private final RealmContext realmContext = () -> "realmName";
5860

5961
private TaskFileIOSupplier buildTaskFileIOSupplier(FileIO fileIO) {
6062
return new TaskFileIOSupplier(new TestFileIOFactory(fileIO));
6163
}
6264

65+
private PolarisCallContext newCallContext() {
66+
BasePersistence metaStore = metaStoreManagerFactory.getOrCreateSession(realmContext);
67+
return new PolarisCallContext(realmContext, metaStore, diagnostics);
68+
}
69+
6370
@Test
6471
public void testMetadataFileCleanup() throws IOException {
65-
PolarisCallContext polarisCallContext =
66-
new PolarisCallContext(
67-
realmContext,
68-
metaStoreManagerFactory.getOrCreateSession(realmContext),
69-
new PolarisDefaultDiagServiceImpl());
72+
PolarisCallContext polarisCallContext = newCallContext();
7073
FileIO fileIO =
7174
new InMemoryFileIO() {
7275
@Override
@@ -175,11 +178,7 @@ public void close() {
175178

176179
@Test
177180
public void testMetadataFileCleanupIfFileNotExist() throws IOException {
178-
PolarisCallContext polarisCallContext =
179-
new PolarisCallContext(
180-
realmContext,
181-
metaStoreManagerFactory.getOrCreateSession(realmContext),
182-
new PolarisDefaultDiagServiceImpl());
181+
PolarisCallContext polarisCallContext = newCallContext();
183182
FileIO fileIO = new InMemoryFileIO();
184183
TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1");
185184
BatchFileCleanupTaskHandler handler =
@@ -219,11 +218,7 @@ public void testMetadataFileCleanupIfFileNotExist() throws IOException {
219218

220219
@Test
221220
public void testCleanupWithRetries() throws IOException {
222-
PolarisCallContext polarisCallContext =
223-
new PolarisCallContext(
224-
realmContext,
225-
metaStoreManagerFactory.getOrCreateSession(realmContext),
226-
new PolarisDefaultDiagServiceImpl());
221+
PolarisCallContext polarisCallContext = newCallContext();
227222
Map<String, AtomicInteger> retryCounter = new HashMap<>();
228223
FileIO fileIO =
229224
new InMemoryFileIO() {

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

Lines changed: 12 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -41,16 +41,18 @@
4141
import org.apache.iceberg.io.OutputFile;
4242
import org.apache.iceberg.io.PositionOutputStream;
4343
import org.apache.polaris.core.PolarisCallContext;
44-
import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
44+
import org.apache.polaris.core.PolarisDiagnostics;
4545
import org.apache.polaris.core.context.RealmContext;
4646
import org.apache.polaris.core.entity.AsyncTaskType;
4747
import org.apache.polaris.core.entity.TaskEntity;
48+
import org.apache.polaris.core.persistence.BasePersistence;
4849
import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
4950
import org.apache.polaris.service.TestFileIOFactory;
5051
import org.junit.jupiter.api.Test;
5152

5253
@QuarkusTest
5354
class ManifestFileCleanupTaskHandlerTest {
55+
@Inject PolarisDiagnostics diagnostics;
5456
@Inject MetaStoreManagerFactory metaStoreManagerFactory;
5557

5658
private final RealmContext realmContext = () -> "realmName";
@@ -59,13 +61,14 @@ private TaskFileIOSupplier buildTaskFileIOSupplier(FileIO fileIO) {
5961
return new TaskFileIOSupplier(new TestFileIOFactory(fileIO));
6062
}
6163

64+
private PolarisCallContext newCallContext() {
65+
BasePersistence metaStore = metaStoreManagerFactory.getOrCreateSession(realmContext);
66+
return new PolarisCallContext(realmContext, metaStore, diagnostics);
67+
}
68+
6269
@Test
6370
public void testCleanupFileNotExists() throws IOException {
64-
PolarisCallContext polarisCallContext =
65-
new PolarisCallContext(
66-
realmContext,
67-
metaStoreManagerFactory.getOrCreateSession(realmContext),
68-
new PolarisDefaultDiagServiceImpl());
71+
PolarisCallContext polarisCallContext = newCallContext();
6972
FileIO fileIO = new InMemoryFileIO();
7073
TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1");
7174

@@ -91,11 +94,7 @@ public void testCleanupFileNotExists() throws IOException {
9194

9295
@Test
9396
public void testCleanupFileManifestExistsDataFilesDontExist() throws IOException {
94-
PolarisCallContext polarisCallContext =
95-
new PolarisCallContext(
96-
realmContext,
97-
metaStoreManagerFactory.getOrCreateSession(realmContext),
98-
new PolarisDefaultDiagServiceImpl());
97+
PolarisCallContext polarisCallContext = newCallContext();
9998
FileIO fileIO = new InMemoryFileIO();
10099
TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1");
101100
ManifestFileCleanupTaskHandler handler =
@@ -119,11 +118,7 @@ public void testCleanupFileManifestExistsDataFilesDontExist() throws IOException
119118

120119
@Test
121120
public void testCleanupFiles() throws IOException {
122-
PolarisCallContext polarisCallContext =
123-
new PolarisCallContext(
124-
realmContext,
125-
metaStoreManagerFactory.getOrCreateSession(realmContext),
126-
new PolarisDefaultDiagServiceImpl());
121+
PolarisCallContext polarisCallContext = newCallContext();
127122
FileIO fileIO =
128123
new InMemoryFileIO() {
129124
@Override
@@ -164,11 +159,7 @@ public void close() {
164159

165160
@Test
166161
public void testCleanupFilesWithRetries() throws IOException {
167-
PolarisCallContext polarisCallContext =
168-
new PolarisCallContext(
169-
realmContext,
170-
metaStoreManagerFactory.getOrCreateSession(realmContext),
171-
new PolarisDefaultDiagServiceImpl());
162+
PolarisCallContext polarisCallContext = newCallContext();
172163
Map<String, AtomicInteger> retryCounter = new HashMap<>();
173164
FileIO fileIO =
174165
new InMemoryFileIO() {

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222
import org.apache.polaris.core.context.CallContext;
2323
import org.apache.polaris.core.context.RealmContext;
2424
import org.apache.polaris.core.entity.TaskEntity;
25-
import org.apache.polaris.core.persistence.BasePersistence;
2625
import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
2726
import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
2827
import org.apache.polaris.service.TestServices;
@@ -47,10 +46,8 @@ void testEventsAreEmitted() {
4746
MetaStoreManagerFactory metaStoreManagerFactory = testServices.metaStoreManagerFactory();
4847
PolarisMetaStoreManager metaStoreManager =
4948
metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext);
50-
BasePersistence bp = metaStoreManagerFactory.getOrCreateSession(realmContext);
5149

52-
PolarisCallContext polarisCallCtx =
53-
new PolarisCallContext(realmContext, bp, testServices.polarisDiagnostics());
50+
PolarisCallContext polarisCallCtx = testServices.newCallContext();
5451

5552
// This task doesn't have a type so it won't be handle-able by a real handler. We register a
5653
// test TaskHandler below that can handle any task.

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -298,4 +298,10 @@ public String getAuthenticationScheme() {
298298
polarisEventListener);
299299
}
300300
}
301+
302+
public PolarisCallContext newCallContext() {
303+
BasePersistence metaStore = metaStoreManagerFactory.getOrCreateSession(realmContext);
304+
return new PolarisCallContext(
305+
realmContext, metaStore, polarisDiagnostics, configurationStore, Clock.systemUTC());
306+
}
301307
}

0 commit comments

Comments
 (0)