Skip to content

Commit b7771df

Browse files
XN137XN137
committed
Inject DefaultFileIOFactory in tests
also simplify `TaskFileIOSupplier` usage in tests, which allows removal of `TestFileIOFactory`.
1 parent 406fa03 commit b7771df

File tree

9 files changed

+65
-117
lines changed

9 files changed

+65
-117
lines changed

‎runtime/service/src/test/java/org/apache/polaris/service/catalog/generic/AbstractPolarisGenericTableCatalogTest.java‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,6 @@
6464
import org.apache.polaris.service.admin.PolarisAdminService;
6565
import org.apache.polaris.service.catalog.PolarisPassthroughResolutionView;
6666
import org.apache.polaris.service.catalog.iceberg.IcebergCatalog;
67-
import org.apache.polaris.service.catalog.io.DefaultFileIOFactory;
6867
import org.apache.polaris.service.catalog.io.FileIOFactory;
6968
import org.apache.polaris.service.catalog.io.StorageAccessConfigProvider;
7069
import org.apache.polaris.service.config.ReservedProperties;
@@ -106,14 +105,14 @@ public abstract class AbstractPolarisGenericTableCatalogTest {
106105
@Inject CallContext callContext;
107106
@Inject RealmConfig realmConfig;
108107
@Inject StorageAccessConfigProvider storageAccessConfigProvider;
108+
@Inject FileIOFactory fileIOFactory;
109109

110110
private PolarisGenericTableCatalog genericTableCatalog;
111111
private IcebergCatalog icebergCatalog;
112112
private AwsStorageConfigInfo storageConfigModel;
113113
private String realmName;
114114
private PolarisCallContext polarisContext;
115115
private PolarisAdminService adminService;
116-
private FileIOFactory fileIOFactory;
117116
private PolarisPrincipal authenticatedRoot;
118117
private PolarisEntity catalogEntity;
119118

@@ -195,7 +194,6 @@ public void before(TestInfo testInfo) {
195194
new PolarisPassthroughResolutionView(
196195
resolutionManifestFactory, authenticatedRoot, CATALOG_NAME);
197196
TaskExecutor taskExecutor = Mockito.mock();
198-
this.fileIOFactory = new DefaultFileIOFactory();
199197

200198
StsClient stsClient = Mockito.mock(StsClient.class);
201199
when(stsClient.assumeRole(isA(AssumeRoleRequest.class)))

‎runtime/service/src/test/java/org/apache/polaris/service/catalog/iceberg/AbstractIcebergCatalogTest.java‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,6 @@
134134
import org.apache.polaris.service.admin.PolarisAdminService;
135135
import org.apache.polaris.service.catalog.PolarisPassthroughResolutionView;
136136
import org.apache.polaris.service.catalog.Profiles;
137-
import org.apache.polaris.service.catalog.io.DefaultFileIOFactory;
138137
import org.apache.polaris.service.catalog.io.ExceptionMappingFileIO;
139138
import org.apache.polaris.service.catalog.io.FileIOFactory;
140139
import org.apache.polaris.service.catalog.io.MeasuredFileIOFactory;
@@ -238,13 +237,14 @@ public Map<String, String> getConfigOverrides() {
238237
@Inject RealmConfig realmConfig;
239238
@Inject ResolutionManifestFactory resolutionManifestFactory;
240239
@Inject StorageAccessConfigProvider storageAccessConfigProvider;
240+
@Inject FileIOFactory fileIOFactory;
241+
@Inject TaskFileIOSupplier taskFileIOSupplier;
241242

242243
private IcebergCatalog catalog;
243244
private String realmName;
244245
private PolarisCallContext polarisContext;
245246
private PolarisAdminService adminService;
246247
private ResolverFactory resolverFactory;
247-
private FileIOFactory fileIOFactory;
248248
private InMemoryFileIO fileIO;
249249
private PolarisEntity catalogEntity;
250250
private PolarisPrincipal authenticatedRoot;
@@ -338,8 +338,6 @@ public void before(TestInfo testInfo) {
338338
.build()
339339
.asCatalog(serviceIdentityProvider)));
340340

341-
this.fileIOFactory = new DefaultFileIOFactory();
342-
343341
StsClient stsClient = Mockito.mock(StsClient.class);
344342
when(stsClient.assumeRole(isA(AssumeRoleRequest.class)))
345343
.thenReturn(
@@ -983,7 +981,7 @@ public void testValidateNotificationFailToCreateFileIO() {
983981
// filename.
984982
final String tableLocation = "s3://externally-owned-bucket/validate_table/";
985983
final String tableMetadataLocation = tableLocation + "metadata/";
986-
FileIOFactory fileIOFactory = spy(new DefaultFileIOFactory());
984+
FileIOFactory fileIOFactory = spy(this.fileIOFactory);
987985
IcebergCatalog catalog = newIcebergCatalog(catalog().name(), metaStoreManager, fileIOFactory);
988986
catalog.initialize(
989987
CATALOG_NAME,
@@ -1899,9 +1897,7 @@ public void testDropTableWithPurge() {
18991897
.containsEntry(StorageAccessProperty.AWS_KEY_ID.getPropertyName(), TEST_ACCESS_KEY)
19001898
.containsEntry(StorageAccessProperty.AWS_SECRET_KEY.getPropertyName(), SECRET_ACCESS_KEY)
19011899
.containsEntry(StorageAccessProperty.AWS_TOKEN.getPropertyName(), SESSION_TOKEN);
1902-
FileIO fileIO =
1903-
new TaskFileIOSupplier(new DefaultFileIOFactory(), storageAccessConfigProvider)
1904-
.apply(taskEntity, TABLE);
1900+
FileIO fileIO = taskFileIOSupplier.apply(taskEntity, TABLE);
19051901
Assertions.assertThat(fileIO).isNotNull().isInstanceOf(ExceptionMappingFileIO.class);
19061902
Assertions.assertThat(((ExceptionMappingFileIO) fileIO).getInnerIo())
19071903
.isInstanceOf(InMemoryFileIO.class);

‎runtime/service/src/test/java/org/apache/polaris/service/catalog/iceberg/AbstractIcebergCatalogViewTest.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,6 @@
5555
import org.apache.polaris.service.admin.PolarisAdminService;
5656
import org.apache.polaris.service.catalog.PolarisPassthroughResolutionView;
5757
import org.apache.polaris.service.catalog.Profiles;
58-
import org.apache.polaris.service.catalog.io.DefaultFileIOFactory;
5958
import org.apache.polaris.service.catalog.io.FileIOFactory;
6059
import org.apache.polaris.service.catalog.io.StorageAccessConfigProvider;
6160
import org.apache.polaris.service.config.ReservedProperties;
@@ -113,6 +112,7 @@ public Map<String, String> getConfigOverrides() {
113112
@Inject CallContext callContext;
114113
@Inject RealmConfig realmConfig;
115114
@Inject StorageAccessConfigProvider storageAccessConfigProvider;
115+
@Inject FileIOFactory fileIOFactory;
116116

117117
private IcebergCatalog catalog;
118118

@@ -189,7 +189,6 @@ public void before(TestInfo testInfo) {
189189
PolarisPassthroughResolutionView passthroughView =
190190
new PolarisPassthroughResolutionView(
191191
resolutionManifestFactory, authenticatedRoot, CATALOG_NAME);
192-
FileIOFactory fileIOFactory = new DefaultFileIOFactory();
193192

194193
testPolarisEventListener = (TestPolarisEventListener) polarisEventListener;
195194
testPolarisEventListener.clear();

‎runtime/service/src/test/java/org/apache/polaris/service/catalog/iceberg/IcebergCatalogHandlerAuthzTest.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,6 @@
6969
import org.apache.polaris.core.persistence.dao.entity.CreatePrincipalResult;
7070
import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifest;
7171
import org.apache.polaris.service.admin.PolarisAuthzTestBase;
72-
import org.apache.polaris.service.catalog.io.DefaultFileIOFactory;
7372
import org.apache.polaris.service.context.catalog.CallContextCatalogFactory;
7473
import org.apache.polaris.service.context.catalog.PolarisCallContextCatalogFactory;
7574
import org.apache.polaris.service.http.IfNoneMatch;
@@ -1891,7 +1890,7 @@ public void testSendNotificationSufficientPrivileges() {
18911890
resolverFactory,
18921891
Mockito.mock(),
18931892
storageAccessConfigProvider,
1894-
new DefaultFileIOFactory(),
1893+
fileIOFactory,
18951894
polarisEventListener,
18961895
metaStoreManager,
18971896
callContext,

‎runtime/service/src/test/java/org/apache/polaris/service/catalog/policy/AbstractPolicyCatalogTest.java‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,6 @@
7676
import org.apache.polaris.service.admin.PolarisAdminService;
7777
import org.apache.polaris.service.catalog.PolarisPassthroughResolutionView;
7878
import org.apache.polaris.service.catalog.iceberg.IcebergCatalog;
79-
import org.apache.polaris.service.catalog.io.DefaultFileIOFactory;
8079
import org.apache.polaris.service.catalog.io.FileIOFactory;
8180
import org.apache.polaris.service.catalog.io.StorageAccessConfigProvider;
8281
import org.apache.polaris.service.config.ReservedProperties;
@@ -132,14 +131,14 @@ public abstract class AbstractPolicyCatalogTest {
132131
@Inject CallContext callContext;
133132
@Inject RealmConfig realmConfig;
134133
@Inject StorageAccessConfigProvider storageAccessConfigProvider;
134+
@Inject FileIOFactory fileIOFactory;
135135

136136
private PolicyCatalog policyCatalog;
137137
private IcebergCatalog icebergCatalog;
138138
private AwsStorageConfigInfo storageConfigModel;
139139
private String realmName;
140140
private PolarisCallContext polarisContext;
141141
private PolarisAdminService adminService;
142-
private FileIOFactory fileIOFactory;
143142
private PolarisPrincipal authenticatedRoot;
144143
private PolarisEntity catalogEntity;
145144

@@ -214,7 +213,6 @@ public void before(TestInfo testInfo) {
214213
new PolarisPassthroughResolutionView(
215214
resolutionManifestFactory, authenticatedRoot, CATALOG_NAME);
216215
TaskExecutor taskExecutor = Mockito.mock();
217-
this.fileIOFactory = new DefaultFileIOFactory();
218216

219217
StsClient stsClient = Mockito.mock(StsClient.class);
220218
when(stsClient.assumeRole(isA(AssumeRoleRequest.class)))

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

Lines changed: 27 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
import static org.assertj.core.api.Assertions.assertThat;
2323
import static org.assertj.core.api.Assertions.assertThatPredicate;
2424

25+
import io.quarkus.test.InjectMock;
26+
import io.quarkus.test.junit.QuarkusMock;
2527
import io.quarkus.test.junit.QuarkusTest;
2628
import jakarta.inject.Inject;
2729
import java.io.IOException;
@@ -30,6 +32,7 @@
3032
import java.util.Map;
3133
import java.util.UUID;
3234
import java.util.concurrent.CompletableFuture;
35+
import java.util.concurrent.ExecutorService;
3336
import java.util.concurrent.Executors;
3437
import java.util.concurrent.atomic.AtomicInteger;
3538
import java.util.stream.Stream;
@@ -43,34 +46,43 @@
4346
import org.apache.iceberg.inmemory.InMemoryFileIO;
4447
import org.apache.iceberg.io.FileIO;
4548
import org.apache.polaris.core.PolarisCallContext;
49+
import org.apache.polaris.core.context.CallContext;
4650
import org.apache.polaris.core.context.RealmContext;
4751
import org.apache.polaris.core.entity.AsyncTaskType;
4852
import org.apache.polaris.core.entity.TaskEntity;
49-
import org.apache.polaris.core.persistence.BasePersistence;
50-
import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
51-
import org.apache.polaris.service.TestFileIOFactory;
52-
import org.apache.polaris.service.catalog.io.StorageAccessConfigProvider;
53+
import org.junit.jupiter.api.AfterEach;
54+
import org.junit.jupiter.api.BeforeEach;
5355
import org.junit.jupiter.api.Test;
5456
import org.mockito.Mockito;
5557

5658
@QuarkusTest
5759
public class BatchFileCleanupTaskHandlerTest {
58-
@Inject MetaStoreManagerFactory metaStoreManagerFactory;
60+
@Inject CallContext callContext;
61+
@InjectMock TaskFileIOSupplier taskFileIOSupplier;
62+
5963
private final RealmContext realmContext = () -> "realmName";
64+
private PolarisCallContext polarisCallContext;
65+
private ExecutorService executor;
66+
67+
@BeforeEach
68+
public void beforeEach() {
69+
QuarkusMock.installMockForType(realmContext, RealmContext.class);
70+
polarisCallContext = callContext.getPolarisCallContext();
71+
executor = Executors.newSingleThreadExecutor();
72+
}
6073

61-
private TaskFileIOSupplier buildTaskFileIOSupplier(FileIO fileIO) {
62-
return new TaskFileIOSupplier(
63-
new TestFileIOFactory(fileIO), Mockito.mock(StorageAccessConfigProvider.class));
74+
@AfterEach
75+
public void afterEach() {
76+
executor.shutdownNow();
6477
}
6578

66-
private PolarisCallContext newCallContext() {
67-
BasePersistence metaStore = metaStoreManagerFactory.getOrCreateSession(realmContext);
68-
return new PolarisCallContext(realmContext, metaStore);
79+
private BatchFileCleanupTaskHandler newBatchFileCleanupTaskHandler(FileIO fileIO) {
80+
Mockito.when(taskFileIOSupplier.apply(Mockito.any(), Mockito.any())).thenReturn(fileIO);
81+
return new BatchFileCleanupTaskHandler(taskFileIOSupplier, executor);
6982
}
7083

7184
@Test
7285
public void testMetadataFileCleanup() throws IOException {
73-
PolarisCallContext polarisCallContext = newCallContext();
7486
FileIO fileIO =
7587
new InMemoryFileIO() {
7688
@Override
@@ -79,9 +91,7 @@ public void close() {
7991
}
8092
};
8193
TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1");
82-
BatchFileCleanupTaskHandler handler =
83-
new BatchFileCleanupTaskHandler(
84-
buildTaskFileIOSupplier(fileIO), Executors.newSingleThreadExecutor());
94+
BatchFileCleanupTaskHandler handler = newBatchFileCleanupTaskHandler(fileIO);
8595

8696
long snapshotId1 = 100L;
8797
ManifestFile manifestFile1 =
@@ -179,12 +189,9 @@ public void close() {
179189

180190
@Test
181191
public void testMetadataFileCleanupIfFileNotExist() throws IOException {
182-
PolarisCallContext polarisCallContext = newCallContext();
183192
FileIO fileIO = new InMemoryFileIO();
184193
TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1");
185-
BatchFileCleanupTaskHandler handler =
186-
new BatchFileCleanupTaskHandler(
187-
buildTaskFileIOSupplier(fileIO), Executors.newSingleThreadExecutor());
194+
BatchFileCleanupTaskHandler handler = newBatchFileCleanupTaskHandler(fileIO);
188195
long snapshotId = 100L;
189196
ManifestFile manifestFile =
190197
TaskTestUtils.manifestFile(
@@ -219,7 +226,6 @@ public void testMetadataFileCleanupIfFileNotExist() throws IOException {
219226

220227
@Test
221228
public void testCleanupWithRetries() throws IOException {
222-
PolarisCallContext polarisCallContext = newCallContext();
223229
Map<String, AtomicInteger> retryCounter = new HashMap<>();
224230
FileIO fileIO =
225231
new InMemoryFileIO() {
@@ -240,9 +246,7 @@ public void deleteFile(String location) {
240246
}
241247
};
242248
TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1");
243-
BatchFileCleanupTaskHandler handler =
244-
new BatchFileCleanupTaskHandler(
245-
buildTaskFileIOSupplier(fileIO), Executors.newSingleThreadExecutor());
249+
BatchFileCleanupTaskHandler handler = newBatchFileCleanupTaskHandler(fileIO);
246250
long snapshotId = 100L;
247251
ManifestFile manifestFile =
248252
TaskTestUtils.manifestFile(

0 commit comments

Comments
 (0)