Skip to content

Commit ef131e2

Browse files
XJDKCXJDKC
committed
Add a test for DefaultPolarisCredentialManager
1 parent d2083a4 commit ef131e2

File tree

3 files changed

+149
-2
lines changed

3 files changed

+149
-2
lines changed

‎polaris-core/src/main/java/org/apache/polaris/core/credentials/DefaultPolarisCredentialManager.java

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
package org.apache.polaris.core.credentials;
2121

22+
import com.google.common.annotations.VisibleForTesting;
2223
import java.util.EnumMap;
2324
import java.util.Optional;
2425
import org.apache.polaris.core.connection.AuthenticationParametersDpo;
@@ -29,6 +30,7 @@
2930
import org.apache.polaris.core.identity.resolved.ResolvedAwsIamServiceIdentity;
3031
import org.apache.polaris.core.identity.resolved.ResolvedServiceIdentity;
3132
import org.slf4j.LoggerFactory;
33+
import software.amazon.awssdk.annotations.NotNull;
3234
import software.amazon.awssdk.services.sts.StsClient;
3335
import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
3436
import software.amazon.awssdk.services.sts.model.AssumeRoleResponse;
@@ -58,7 +60,7 @@ public EnumMap<ConnectionCredentialProperty, String> getConnectionCredentials(
5860
(ResolvedAwsIamServiceIdentity) resolvedServiceIdentity;
5961
SigV4AuthenticationParametersDpo sigV4AuthenticationParameters =
6062
(SigV4AuthenticationParametersDpo) authenticationParameters;
61-
StsClient stsClient = resolvedAwsIamServiceIdentity.stsClientSupplier().get();
63+
StsClient stsClient = getStsClient(resolvedAwsIamServiceIdentity);
6264
AssumeRoleResponse response =
6365
stsClient.assumeRole(
6466
AssumeRoleRequest.builder()
@@ -90,4 +92,10 @@ public EnumMap<ConnectionCredentialProperty, String> getConnectionCredentials(
9092
}
9193
return credentialMap;
9294
}
95+
96+
@VisibleForTesting
97+
public StsClient getStsClient(
98+
@NotNull ResolvedAwsIamServiceIdentity resolvedAwsIamServiceIdentity) {
99+
return resolvedAwsIamServiceIdentity.stsClientSupplier().get();
100+
}
93101
}

‎quarkus/service/src/test/java/org/apache/polaris/service/quarkus/credentials/DefaultPolarisCredentialManagerTest.java

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.polaris.service.quarkus.credentials;
21+
22+
import static org.mockito.ArgumentMatchers.any;
23+
import static org.mockito.Mockito.*;
24+
25+
import io.quarkus.test.InjectMock;
26+
import io.quarkus.test.junit.QuarkusTest;
27+
import io.quarkus.test.junit.QuarkusTestProfile;
28+
import io.quarkus.test.junit.TestProfile;
29+
import jakarta.inject.Inject;
30+
import java.time.Instant;
31+
import java.util.EnumMap;
32+
import java.util.Map;
33+
import org.apache.polaris.core.connection.SigV4AuthenticationParametersDpo;
34+
import org.apache.polaris.core.context.RealmContext;
35+
import org.apache.polaris.core.credentials.DefaultPolarisCredentialManager;
36+
import org.apache.polaris.core.credentials.connection.ConnectionCredentialProperty;
37+
import org.apache.polaris.core.identity.ServiceIdentityType;
38+
import org.apache.polaris.core.identity.dpo.ServiceIdentityInfoDpo;
39+
import org.apache.polaris.core.identity.registry.ServiceIdentityRegistry;
40+
import org.apache.polaris.core.identity.resolved.ResolvedAwsIamServiceIdentity;
41+
import org.assertj.core.api.Assertions;
42+
import org.junit.jupiter.api.BeforeEach;
43+
import org.junit.jupiter.api.Test;
44+
import org.mockito.Mockito;
45+
import software.amazon.awssdk.services.sts.StsClient;
46+
import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
47+
import software.amazon.awssdk.services.sts.model.AssumeRoleResponse;
48+
import software.amazon.awssdk.services.sts.model.Credentials;
49+
50+
@QuarkusTest
51+
@TestProfile(DefaultPolarisCredentialManagerTest.Profile.class)
52+
public class DefaultPolarisCredentialManagerTest {
53+
54+
@InjectMock RealmContext realmContext;
55+
56+
@Inject QuarkusPolarisCredentialManagerConfiguration configuration;
57+
@Inject ServiceIdentityRegistry serviceIdentityRegistry;
58+
59+
DefaultPolarisCredentialManager credentialManager;
60+
61+
public static class Profile implements QuarkusTestProfile {
62+
@Override
63+
public Map<String, String> getConfigOverrides() {
64+
return Map.of(
65+
"quarkus.identity-registry.type",
66+
"default",
67+
"polaris.service-identity.my-realm.aws-iam.iam-arn",
68+
"arn:aws:iam::123456789012:user/polaris-iam-user",
69+
"polaris.service-identity.my-realm.aws-iam.access-key-id",
70+
"access-key-id",
71+
"polaris.service-identity.my-realm.aws-iam.secret-access-key",
72+
"secret-access-key",
73+
"polaris.credential-manager.type",
74+
"default");
75+
}
76+
}
77+
78+
@BeforeEach
79+
void setup() {
80+
// Mock the realm context to return a specific realm
81+
when(realmContext.getRealmIdentifier()).thenReturn("my-realm");
82+
83+
credentialManager = Mockito.spy(new DefaultPolarisCredentialManager(serviceIdentityRegistry));
84+
doAnswer(
85+
invocation -> {
86+
// Capture the identity here
87+
ResolvedAwsIamServiceIdentity identity = invocation.getArgument(0);
88+
89+
StsClient mockStsClient = mock(StsClient.class);
90+
when(mockStsClient.assumeRole(Mockito.any(AssumeRoleRequest.class)))
91+
.thenAnswer(
92+
stsInvocation -> {
93+
// Validate identity at the time assumeRole is called
94+
if (!"access-key-id".equals(identity.getAccessKeyId())
95+
|| !"secret-access-key".equals(identity.getSecretAccessKey())) {
96+
throw new IllegalArgumentException("Invalid credentials on assumeRole");
97+
}
98+
99+
// Return mocked credentials
100+
Credentials tmpSessionCredentials =
101+
Credentials.builder()
102+
.accessKeyId("tmp-access-key-id")
103+
.secretAccessKey("tmp-secret-access-key")
104+
.sessionToken("tmp-session-token")
105+
.expiration(Instant.now().plusSeconds(3600))
106+
.build();
107+
108+
return AssumeRoleResponse.builder()
109+
.credentials(tmpSessionCredentials)
110+
.build();
111+
});
112+
return mockStsClient;
113+
})
114+
.when(credentialManager)
115+
.getStsClient(any());
116+
}
117+
118+
@Test
119+
public void testGetConnectionCredentialsForSigV4() {
120+
ServiceIdentityInfoDpo serviceIdentityInfo =
121+
serviceIdentityRegistry.assignServiceIdentity(ServiceIdentityType.AWS_IAM);
122+
EnumMap<ConnectionCredentialProperty, String> credentials =
123+
credentialManager.getConnectionCredentials(
124+
serviceIdentityInfo,
125+
new SigV4AuthenticationParametersDpo(
126+
"arn:aws:iam::123456789012:role/polaris-users-iam-role",
127+
null,
128+
null,
129+
"us-west-2",
130+
"glue"));
131+
Assertions.assertThat(credentials)
132+
.containsEntry(ConnectionCredentialProperty.AWS_ACCESS_KEY_ID, "tmp-access-key-id")
133+
.containsEntry(ConnectionCredentialProperty.AWS_SECRET_ACCESS_KEY, "tmp-secret-access-key")
134+
.containsEntry(ConnectionCredentialProperty.AWS_SESSION_TOKEN, "tmp-session-token")
135+
.containsKey(ConnectionCredentialProperty.EXPIRATION_TIME)
136+
.size()
137+
.isEqualTo(4);
138+
}
139+
}

‎service/common/src/main/java/org/apache/polaris/service/identity/mutation/EntityMutationConfiguration.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,5 +23,5 @@
2323
import java.util.Optional;
2424

2525
public interface EntityMutationConfiguration {
26-
Optional<List<String>> mutators();
26+
Optional<List<String>> mutators();
2727
}

0 commit comments

Comments
 (0)