Skip to content

Commit f1d20ef

Browse files
adutraadutra
committed
chore(tests): Reduce flakiness due to MockServer websocket errors
Some test flakiness is happening lately, always related to MockServer and how it creates WebSockets. The flakiness seems to only affect macOS, CI runs fine. This PR attempts to mitigate the flakiness by tweaking some parameters heuristically: 1. Lower the number of expectations that require WebSockets 2. Lower the number of test permutations in OAuth2AgentTest 3. Increase some thread pool sizes for MockServer server and client 4. Explicitly clear all expectations before closing the server.
1 parent f4aa9e0 commit f1d20ef

13 files changed

Lines changed: 100 additions & 114 deletions

‎oauth2/core/src/test/java/com/dremio/iceberg/authmgr/oauth2/agent/OAuth2AgentTest.java‎

Lines changed: 8 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
import static org.assertj.core.api.Assertions.assertThatCode;
2222
import static org.assertj.core.api.InstanceOfAssertFactories.ATOMIC_BOOLEAN;
2323
import static org.assertj.core.api.InstanceOfAssertFactories.throwable;
24-
import static org.junit.jupiter.api.Assumptions.assumeTrue;
2524
import static org.mockito.ArgumentMatchers.any;
2625
import static org.mockito.ArgumentMatchers.anyLong;
2726
import static org.mockito.Mockito.doAnswer;
@@ -357,20 +356,15 @@ void testTokenExchangeStaticSubjectActor(
357356
@CartesianTest
358357
void testTokenExchangeDynamicSubject(
359358
@Enum HttpClientType httpClientType,
360-
@EnumLike ClientAuthenticationMethod authenticationMethod,
361-
@Values(booleans = {true, false}) boolean returnRefreshTokens,
362-
@EnumLike(excludes = "urn:ietf:params:oauth:grant-type:token-exchange")
363-
GrantType subjectGrantType)
359+
@EnumLike(excludes = "none") ClientAuthenticationMethod authenticationMethod,
360+
@Values(booleans = {true, false}) boolean returnRefreshTokens)
364361
throws InterruptedException, ExecutionException {
365-
assumeTrue(
366-
!subjectGrantType.equals(GrantType.CLIENT_CREDENTIALS)
367-
|| !authenticationMethod.equals(ClientAuthenticationMethod.NONE));
368362
try (TestEnvironment env =
369363
TestEnvironment.builder()
370364
.grantType(GrantType.TOKEN_EXCHANGE)
371-
.httpClientType(httpClientType)
372365
.subjectToken(null)
373-
.subjectGrantType(subjectGrantType)
366+
.subjectGrantType(GrantType.CLIENT_CREDENTIALS)
367+
.httpClientType(httpClientType)
374368
.clientAuthenticationMethod(authenticationMethod)
375369
.returnRefreshTokens(returnRefreshTokens)
376370
.build();
@@ -387,20 +381,15 @@ void testTokenExchangeDynamicSubject(
387381
@CartesianTest
388382
void testTokenExchangeDynamicActor(
389383
@Enum HttpClientType httpClientType,
390-
@EnumLike ClientAuthenticationMethod authenticationMethod,
391-
@Values(booleans = {true, false}) boolean returnRefreshTokens,
392-
@EnumLike(excludes = "urn:ietf:params:oauth:grant-type:token-exchange")
393-
GrantType actorGrantType)
384+
@EnumLike(excludes = "none") ClientAuthenticationMethod authenticationMethod,
385+
@Values(booleans = {true, false}) boolean returnRefreshTokens)
394386
throws InterruptedException, ExecutionException {
395-
assumeTrue(
396-
!actorGrantType.equals(GrantType.CLIENT_CREDENTIALS)
397-
|| !authenticationMethod.equals(ClientAuthenticationMethod.NONE));
398387
try (TestEnvironment env =
399388
TestEnvironment.builder()
400389
.grantType(GrantType.TOKEN_EXCHANGE)
401-
.httpClientType(httpClientType)
402390
.actorToken(null)
403-
.actorGrantType(actorGrantType)
391+
.actorGrantType(GrantType.CLIENT_CREDENTIALS)
392+
.httpClientType(httpClientType)
404393
.clientAuthenticationMethod(authenticationMethod)
405394
.returnRefreshTokens(returnRefreshTokens)
406395
.build();

‎oauth2/core/src/test/java/com/dremio/iceberg/authmgr/oauth2/flow/TokenExchangeFlowTest.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,8 +49,7 @@ void fetchNewTokens(
4949

5050
@CartesianTest
5151
void fetchNewTokensDynamic(
52-
@EnumLike(excludes = {"none", "client_secret_basic"})
53-
ClientAuthenticationMethod authenticationMethod,
52+
@EnumLike(excludes = "none") ClientAuthenticationMethod authenticationMethod,
5453
@Values(booleans = {true, false}) boolean returnRefreshTokens,
5554
@EnumLike(includes = {"client_credentials", "authorization_code"}) GrantType subjectGrantType,
5655
@EnumLike(includes = {"password", "urn:ietf:params:oauth:grant-type:device_code"})

‎oauth2/core/src/test/resources/logback-test.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,5 +29,5 @@ limitations under the License.
2929
<logger name="com.dremio.iceberg.authmgr.oauth2" level="INFO"/>
3030
<logger name="org.apache.iceberg" level="INFO"/>
3131
<!-- MockServer: set level to INFO to print all requests and responses -->
32-
<logger name="org.mockserver.log" level="WARN"/>
32+
<logger name="org.mockserver" level="WARN"/>
3333
</configuration>

‎oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/TestEnvironment.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,9 @@ public void reset() {
166166
@Override
167167
public void close() {
168168
getUser().close();
169+
if (isUnitTest()) {
170+
reset();
171+
}
169172
try {
170173
getExecutor().shutdown();
171174
if (!getExecutor().awaitTermination(10, TimeUnit.SECONDS)) {

‎oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/expectation/AbstractTokenEndpointExpectation.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,8 +53,7 @@ protected HttpRequest request() {
5353
.withBody(ExpectationUtils.getParameterBody(requestBody().build()));
5454
}
5555

56-
protected HttpResponse response(
57-
HttpRequest httpRequest, String accessToken, String refreshToken) {
56+
protected HttpResponse response(String accessToken, String refreshToken) {
5857
return HttpResponse.response()
5958
.withBody(getJsonBody(responseBody(accessToken, refreshToken).build()));
6059
}

‎oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/expectation/AuthorizationCodeExpectation.java‎

Lines changed: 38 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@
4040

4141
@AuthManagerImmutable
4242
@Value.Enclosing
43-
public abstract class AuthorizationCodeExpectation extends InitialTokenFetchExpectation {
43+
public abstract class AuthorizationCodeExpectation extends AbstractTokenEndpointExpectation {
4444

4545
/** A map of pending authorization requests, keyed by the redirect URI. */
4646
@Value.Lazy
@@ -52,7 +52,7 @@ public abstract class AuthorizationCodeExpectation extends InitialTokenFetchExpe
5252
@Override
5353
public void create() {
5454
createAuthEndpointExpectation();
55-
super.create();
55+
createTokenEndpointExpectation();
5656
}
5757

5858
@Override
@@ -68,39 +68,42 @@ protected ImmutableMap.Builder<String, String> requestBody() {
6868
return builder;
6969
}
7070

71-
@Override
72-
protected HttpResponse response(
73-
HttpRequest httpRequest, String accessToken, String refreshToken) {
74-
Map<String, List<String>> params = decodeBodyParameters(httpRequest);
75-
String redirectUri = params.get("redirect_uri").get(0);
76-
PendingAuthRequest pendingAuthRequest = getPendingAuthRequests().get(redirectUri);
77-
if (pendingAuthRequest == null) {
78-
return AUTHORIZATION_SERVER_ERROR_RESPONSE;
79-
}
80-
List<String> code = params.get("code");
81-
if (code == null
82-
|| code.isEmpty()
83-
|| !code.get(0).equals(pendingAuthRequest.getCode().getValue())) {
84-
return AUTHORIZATION_SERVER_ERROR_RESPONSE;
85-
}
86-
if (getTestEnvironment().isPkceEnabled()) {
87-
if (pendingAuthRequest.getCodeChallengeMethod().isEmpty()
88-
|| pendingAuthRequest.getCodeChallenge().isEmpty()) {
89-
return AUTHORIZATION_SERVER_ERROR_RESPONSE;
90-
}
91-
List<String> codeVerifier = params.get("code_verifier");
92-
if (codeVerifier == null || codeVerifier.isEmpty()) {
93-
return AUTHORIZATION_SERVER_ERROR_RESPONSE;
94-
}
95-
if (!pendingAuthRequest
96-
.getCodeChallenge()
97-
.get()
98-
.equals(pendingAuthRequest.getCodeChallenge().get())) {
99-
return AUTHORIZATION_SERVER_ERROR_RESPONSE;
100-
}
101-
}
102-
getPendingAuthRequests().remove(redirectUri);
103-
return super.response(httpRequest, accessToken, refreshToken);
71+
private void createTokenEndpointExpectation() {
72+
getClientAndServer()
73+
.when(request())
74+
.respond(
75+
request -> {
76+
Map<String, List<String>> params = decodeBodyParameters(request);
77+
String redirectUri = params.get("redirect_uri").get(0);
78+
PendingAuthRequest pendingAuthRequest = getPendingAuthRequests().get(redirectUri);
79+
if (pendingAuthRequest == null) {
80+
return AUTHORIZATION_SERVER_ERROR_RESPONSE;
81+
}
82+
List<String> code = params.get("code");
83+
if (code == null
84+
|| code.isEmpty()
85+
|| !code.get(0).equals(pendingAuthRequest.getCode().getValue())) {
86+
return AUTHORIZATION_SERVER_ERROR_RESPONSE;
87+
}
88+
if (getTestEnvironment().isPkceEnabled()) {
89+
if (pendingAuthRequest.getCodeChallengeMethod().isEmpty()
90+
|| pendingAuthRequest.getCodeChallenge().isEmpty()) {
91+
return AUTHORIZATION_SERVER_ERROR_RESPONSE;
92+
}
93+
List<String> codeVerifier = params.get("code_verifier");
94+
if (codeVerifier == null || codeVerifier.isEmpty()) {
95+
return AUTHORIZATION_SERVER_ERROR_RESPONSE;
96+
}
97+
if (!pendingAuthRequest
98+
.getCodeChallenge()
99+
.get()
100+
.equals(pendingAuthRequest.getCodeChallenge().get())) {
101+
return AUTHORIZATION_SERVER_ERROR_RESPONSE;
102+
}
103+
}
104+
getPendingAuthRequests().remove(redirectUri);
105+
return super.response("access_initial", "refresh_initial");
106+
});
104107
}
105108

106109
private void createAuthEndpointExpectation() {

‎oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/expectation/ClientCredentialsExpectation.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,12 @@
2323
import com.nimbusds.oauth2.sdk.GrantType;
2424

2525
@AuthManagerImmutable
26-
public abstract class ClientCredentialsExpectation extends InitialTokenFetchExpectation {
26+
public abstract class ClientCredentialsExpectation extends AbstractTokenEndpointExpectation {
27+
28+
@Override
29+
public void create() {
30+
getClientAndServer().when(request()).respond(response("access_initial", "refresh_initial"));
31+
}
2732

2833
@Override
2934
protected ImmutableMap.Builder<String, String> requestBody() {

‎oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/expectation/DeviceCodeExpectation.java‎

Lines changed: 22 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@
3939
import org.mockserver.model.StringBody;
4040

4141
@AuthManagerImmutable
42-
public abstract class DeviceCodeExpectation extends InitialTokenFetchExpectation {
42+
public abstract class DeviceCodeExpectation extends AbstractTokenEndpointExpectation {
4343

4444
/** A map of pending authorization requests, keyed by the user and device code. */
4545
@Value.Lazy
@@ -52,7 +52,7 @@ public abstract class DeviceCodeExpectation extends InitialTokenFetchExpectation
5252
public void create() {
5353
createDeviceAuthEndpointExpectation();
5454
createDeviceVerificationEndpointExpectation();
55-
super.create();
55+
createTokenEndpointExpectation();
5656
}
5757

5858
@Override
@@ -62,23 +62,26 @@ protected ImmutableMap.Builder<String, String> requestBody() {
6262
.put("device_code", "[a-zA-Z0-9-._~]+");
6363
}
6464

65-
@Override
66-
protected HttpResponse response(
67-
HttpRequest httpRequest, String accessToken, String refreshToken) {
68-
Map<String, List<String>> params = decodeBodyParameters(httpRequest);
69-
DeviceCode deviceCode = new DeviceCode(params.get("device_code").get(0));
70-
PendingAuthRequest pendingAuthRequest = getPendingAuthRequests().get(deviceCode);
71-
if (pendingAuthRequest.isUserCodeReceived()) {
72-
getPendingAuthRequests().remove(pendingAuthRequest.getDeviceCode());
73-
getPendingAuthRequests().remove(pendingAuthRequest.getUserCode());
74-
return super.response(httpRequest, accessToken, refreshToken);
75-
} else {
76-
return HttpResponse.response()
77-
.withStatusCode(401)
78-
.withBody(
79-
JsonBody.json(
80-
"{\"error\":\"authorization_pending\",\"error_description\":\"User code not yet received\"}"));
81-
}
65+
private void createTokenEndpointExpectation() {
66+
getClientAndServer()
67+
.when(request())
68+
.respond(
69+
request -> {
70+
Map<String, List<String>> params = decodeBodyParameters(request);
71+
DeviceCode deviceCode = new DeviceCode(params.get("device_code").get(0));
72+
PendingAuthRequest pendingAuthRequest = getPendingAuthRequests().get(deviceCode);
73+
if (pendingAuthRequest.isUserCodeReceived()) {
74+
getPendingAuthRequests().remove(pendingAuthRequest.getDeviceCode());
75+
getPendingAuthRequests().remove(pendingAuthRequest.getUserCode());
76+
return super.response("access_initial", "refresh_initial");
77+
} else {
78+
return HttpResponse.response()
79+
.withStatusCode(401)
80+
.withBody(
81+
JsonBody.json(
82+
"{\"error\":\"authorization_pending\",\"error_description\":\"User code not yet received\"}"));
83+
}
84+
});
8285
}
8386

8487
private void createDeviceAuthEndpointExpectation() {

‎oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/expectation/InitialTokenFetchExpectation.java‎

Lines changed: 0 additions & 26 deletions
This file was deleted.

‎oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/expectation/PasswordExpectation.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,12 @@
2424
import com.google.common.collect.ImmutableMap;
2525

2626
@AuthManagerImmutable
27-
public abstract class PasswordExpectation extends InitialTokenFetchExpectation {
27+
public abstract class PasswordExpectation extends AbstractTokenEndpointExpectation {
28+
29+
@Override
30+
public void create() {
31+
getClientAndServer().when(request()).respond(response("access_initial", "refresh_initial"));
32+
}
2833

2934
@Override
3035
protected ImmutableMap.Builder<String, String> requestBody() {

0 commit comments

Comments
 (0)