Skip to content

Commit

Permalink
chore(spanner): fix unit tests for mux read-write (#3598)
Browse files Browse the repository at this point in the history 
* chore(spanner): fix tests in connect api for mux rw

* chore(spanner): remove test skips

* chore(spanner): add env to verify tests

* chore(spanner): fix tests in AsyncRunnerTest file

* chore(spanner): skip RetryOnDifferentGrpcChannelMockServerTest for mux read-write as it is not handled

* chore(spanner): fix DatabaseClientImpl test

* chore(spanner): lint fix

* chore(spanner): fix SpanTest unit test

* chore(spanner): fix OpenTelemetrySpanTest

* chore(spanner): fix AsyncTransactionManagerTest

* chore(spanner): revert skipped tests to another PR

* chore(spanner): revert env changes

* chore(spanner): lint fix
  • Loading branch information
@harshachinta
harshachinta authored Jan 20, 2025
1 parent cc626f8 commit 2d65e88
Show file tree
Hide file tree
Showing 11 changed files with 361 additions and 73 deletions.
2 changes: 2 additions & 0 deletions .github/workflows/ci.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ jobs:
JOB_TYPE: test
GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: true
GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS: true
GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW: true
units-java8:
# Building using Java 17 and run the tests with Java 8 runtime
name: "units (8)"
Expand Down Expand Up @@ -94,6 +95,7 @@ jobs:
JOB_TYPE: test
GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: true
GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS: true
GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW: true
windows:
runs-on: windows-latest
steps:
Expand Down
5 changes: 5 additions & 0 deletions .kokoro/presubmit/integration-multiplexed-sessions-enabled.cfg
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,3 +41,8 @@ env_vars: {
key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS"
value: "true"
}

env_vars: {
key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW"
value: "true"
}
43 changes: 39 additions & 4 deletions google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -201,7 +201,17 @@ 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,
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.
BeginTransactionRequest.class,
ExecuteSqlRequest.class,
CommitRequest.class);
} else if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactly(
CreateSessionRequest.class,
Expand Down Expand Up @@ -260,7 +270,11 @@ public void asyncRunnerWaitsUntilAsyncUpdateHasFinished() throws Exception {
},
executor);
res.get();
if (isMultiplexedSessionsEnabled()) {
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsAtLeast(
CreateSessionRequest.class, ExecuteSqlRequest.class, CommitRequest.class);
} else if (isMultiplexedSessionsEnabled()) {
// 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())
Expand Down Expand Up @@ -404,7 +418,17 @@ 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,
ExecuteSqlRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class,
ExecuteSqlRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
} else if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactly(
CreateSessionRequest.class,
Expand Down Expand Up @@ -463,7 +487,11 @@ public void asyncRunnerWaitsUntilAsyncBatchUpdateHasFinished() throws Exception
},
executor);
res.get();
if (isMultiplexedSessionsEnabled()) {
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactly(
CreateSessionRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
} else if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactly(
CreateSessionRequest.class,
Expand Down Expand Up @@ -576,4 +604,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();
}
}
189 changes: 153 additions & 36 deletions google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -179,9 +179,11 @@ public void asyncTransactionManager_shouldRollbackOnCloseAsync() throws Exceptio
AsyncTransactionManager manager = client().transactionManagerAsync();
TransactionContext txn = manager.beginAsync().get();
txn.executeUpdateAsync(UPDATE_STATEMENT).get();
final TransactionSelector selector =
((TransactionContextImpl) ((SessionPoolTransactionContext) txn).delegate)
.getTransactionSelector();
if (txn instanceof SessionPoolTransactionContext) {
txn = ((SessionPoolTransactionContext) txn).delegate;
}
TransactionContextImpl impl = (TransactionContextImpl) txn;
final TransactionSelector selector = impl.getTransactionSelector();

SpannerApiFutures.get(manager.closeAsync());
// The mock server should already have the Rollback request, as we are waiting for the returned
Expand Down Expand Up @@ -359,7 +361,22 @@ public void asyncTransactionManagerFireAndForgetInvalidUpdate() throws Exception
ExecuteSqlRequest.class,
ExecuteSqlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionForRW =
ImmutableList.of(
CreateSessionRequest.class,
// The first update that fails. This will cause a transaction retry.
ExecuteSqlRequest.class,
// The retry will use an explicit BeginTransaction call.
BeginTransactionRequest.class,
// The first update will again fail, but now there is a transaction id, so the
// transaction can continue.
ExecuteSqlRequest.class,
ExecuteSqlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionForRW);
} else if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
Expand Down Expand Up @@ -502,14 +519,25 @@ public void asyncTransactionManagerUpdateAbortedWithoutGettingResult() throws Ex
// The server may receive 1 or 2 commit requests depending on whether the call to
// commitAsync() already knows that the transaction has aborted. If it does, it will not
// attempt to call the Commit RPC and instead directly propagate the Aborted error.
assertThat(mockSpanner.getRequestTypes())
.containsAtLeast(
BatchCreateSessionsRequest.class,
ExecuteSqlRequest.class,
// The retry will use a BeginTransaction RPC.
BeginTransactionRequest.class,
ExecuteSqlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsAtLeast(
CreateSessionRequest.class,
ExecuteSqlRequest.class,
// The retry will use a BeginTransaction RPC.
BeginTransactionRequest.class,
ExecuteSqlRequest.class,
CommitRequest.class);
} else {
assertThat(mockSpanner.getRequestTypes())
.containsAtLeast(
BatchCreateSessionsRequest.class,
ExecuteSqlRequest.class,
// The retry will use a BeginTransaction RPC.
BeginTransactionRequest.class,
ExecuteSqlRequest.class,
CommitRequest.class);
}
break;
} catch (AbortedException e) {
transactionContextFuture = manager.resetForRetryAsync();
Expand Down Expand Up @@ -557,7 +585,11 @@ public void asyncTransactionManagerWaitsUntilAsyncUpdateHasFinished() throws Exc
executor)
.commitAsync()
.get();
if (isMultiplexedSessionsEnabled()) {
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactly(
CreateSessionRequest.class, ExecuteSqlRequest.class, CommitRequest.class);
} else if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactly(
CreateSessionRequest.class,
Expand Down Expand Up @@ -678,7 +710,16 @@ public void asyncTransactionManagerFireAndForgetInvalidBatchUpdate() throws Exce
ExecuteBatchDmlRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionsRW =
ImmutableList.of(
CreateSessionRequest.class,
ExecuteBatchDmlRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW);
} else if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
Expand Down Expand Up @@ -722,7 +763,17 @@ public void asyncTransactionManagerBatchUpdateAborted() throws Exception {
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionsRW =
ImmutableList.of(
CreateSessionRequest.class,
ExecuteBatchDmlRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW);
} else if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
Expand Down Expand Up @@ -764,7 +815,20 @@ public void asyncTransactionManagerBatchUpdateAbortedBeforeFirstStatement() thro
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
// When requests run using multiplexed session with read-write enabled, the
// BatchCreateSessionsRequest will not be
// triggered because we are creating an empty pool during initialization.
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionsRW =
ImmutableList.of(
CreateSessionRequest.class,
ExecuteBatchDmlRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW);
} else if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
Expand Down Expand Up @@ -825,7 +889,18 @@ public void asyncTransactionManagerWithBatchUpdateCommitAborted() throws Excepti
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionsRW =
ImmutableList.of(
CreateSessionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW);
} else if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
Expand Down Expand Up @@ -867,27 +942,50 @@ public void asyncTransactionManagerBatchUpdateAbortedWithoutGettingResult() thro
assertThat(attempt.get()).isEqualTo(2);
List<Class<? extends AbstractMessage>> requests = mockSpanner.getRequestTypes();
// Remove the CreateSession requests for multiplexed sessions, as those are not relevant for
// this test.
requests.removeIf(request -> request == CreateSessionRequest.class);
// this test if multiplexed session for read-write is not enabled.
if (!isMultiplexedSessionsEnabledForRW()) {
requests.removeIf(request -> request == CreateSessionRequest.class);
}
int size = Iterables.size(requests);
assertThat(size).isIn(Range.closed(5, 6));
if (size == 5) {
assertThat(requests)
.containsExactly(
BatchCreateSessionsRequest.class,
ExecuteBatchDmlRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(requests)
.containsExactly(
CreateSessionRequest.class,
ExecuteBatchDmlRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
} else {
assertThat(requests)
.containsExactly(
BatchCreateSessionsRequest.class,
ExecuteBatchDmlRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
}
} else {
assertThat(requests)
.containsExactly(
BatchCreateSessionsRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(requests)
.containsExactly(
CreateSessionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
} else {
assertThat(requests)
.containsExactly(
BatchCreateSessionsRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
}
}
}

Expand Down Expand Up @@ -918,7 +1016,13 @@ public void asyncTransactionManagerWithBatchUpdateCommitFails() {
ImmutableList<Class<? extends Message>> expectedRequests =
ImmutableList.of(
BatchCreateSessionsRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionsRW =
ImmutableList.of(
CreateSessionRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW);
} else if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
Expand Down Expand Up @@ -949,7 +1053,13 @@ public void asyncTransactionManagerWaitsUntilAsyncBatchUpdateHasFinished() throw
ImmutableList<Class<? extends Message>> expectedRequests =
ImmutableList.of(
BatchCreateSessionsRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionsRW =
ImmutableList.of(
CreateSessionRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW);
} else if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
Expand Down Expand Up @@ -1122,4 +1232,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();
}
}
Loading

0 comments on commit 2d65e88

Please sign in to comment.