Skip to content

Commit

Permalink
chore(spanner): fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
harshachinta committed Dec 18, 2024
1 parent dfd3812 commit 3b67313
Showing 1 changed file with 12 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -201,11 +201,10 @@ public void asyncRunnerUpdateAbortedWithoutGettingResult() throws Exception {
executor);
assertThat(result.get()).isNull();
assertThat(attempt.get()).isEqualTo(2);
if (isMultiplexedSessionsEnabled()) {
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactly(
CreateSessionRequest.class,
BatchCreateSessionsRequest.class,
ExecuteSqlRequest.class,
// The retry will use an explicit BeginTransaction RPC because the first statement of
// the transaction did not return a transaction id during the initial attempt.
Expand Down Expand Up @@ -260,12 +259,12 @@ public void asyncRunnerWaitsUntilAsyncUpdateHasFinished() throws Exception {
},
executor);
res.get();
if (isMultiplexedSessionsEnabled()) {
if (isMultiplexedSessionsEnabledForRW()) {
// The mock server could have received a CreateSession request for a multiplexed session, but
// it could also be that that request has not yet reached the server.
assertThat(mockSpanner.getRequestTypes())
.containsAtLeast(
BatchCreateSessionsRequest.class, ExecuteSqlRequest.class, CommitRequest.class);
CreateSessionRequest.class, ExecuteSqlRequest.class, CommitRequest.class);
} else {
assertThat(mockSpanner.getRequestTypes())
.containsExactly(
Expand Down Expand Up @@ -404,11 +403,10 @@ public void asyncRunnerBatchUpdateAbortedWithoutGettingResult() throws Exception
executor);
assertThat(result.get()).isNull();
assertThat(attempt.get()).isEqualTo(2);
if (isMultiplexedSessionsEnabled()) {
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactly(
CreateSessionRequest.class,
BatchCreateSessionsRequest.class,
ExecuteSqlRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class,
Expand Down Expand Up @@ -463,11 +461,10 @@ public void asyncRunnerWaitsUntilAsyncBatchUpdateHasFinished() throws Exception
},
executor);
res.get();
if (isMultiplexedSessionsEnabled()) {
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactly(
CreateSessionRequest.class,
BatchCreateSessionsRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
} else {
Expand Down Expand Up @@ -576,4 +573,11 @@ private boolean isMultiplexedSessionsEnabled() {
}
return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession();
}

private boolean isMultiplexedSessionsEnabledForRW() {
if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) {
return false;
}
return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
}
}

0 comments on commit 3b67313

Please sign in to comment.