Skip to content

Commit 583d3b9

Browse files
amanda-tarafaamanda-tarafa
committed
tests(Spanner): Explicit ManagedSession and ManagedTransaction unit tests
1 parent 2df72c7 commit 583d3b9

File tree

3 files changed

+673
-6
lines changed

3 files changed

+673
-6
lines changed

‎apis/Google.Cloud.Spanner.Data/Google.Cloud.Spanner.Data.Tests/V1/ManagedSessionTests.cs‎

Lines changed: 345 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,345 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License"):
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
using Google.Api.Gax.Grpc;
16+
using Google.Api.Gax.Testing;
17+
using Google.Cloud.Spanner.Common.V1;
18+
using Google.Cloud.Spanner.V1.Internal.Logging;
19+
using Google.Protobuf.WellKnownTypes;
20+
using NSubstitute;
21+
using NSubstitute.Extensions;
22+
using System;
23+
using System.Threading;
24+
using System.Threading.Tasks;
25+
using Xunit;
26+
27+
namespace Google.Cloud.Spanner.V1.Tests;
28+
29+
public class ManagedSessionTests
30+
{
31+
private static readonly DatabaseName SampleDatabaseName = new DatabaseName("project", "instance", "database");
32+
private static readonly TimeSpan SoftExpiry = TimeSpan.FromDays(7);
33+
private static readonly TimeSpan HardExpiry = TimeSpan.FromDays(27) + TimeSpan.FromHours(12);
34+
35+
[Fact]
36+
public async Task GetFreshSession_InitialCreation()
37+
{
38+
var clock = new FakeClock(new DateTime(2025, 1, 1, 0, 0, 0, DateTimeKind.Utc));
39+
var client = SpannerClientHelpers.CreateMockClient(Logger.DefaultLogger, clock);
40+
41+
var options = ManagedSessionOptions.Create(SampleDatabaseName, client);
42+
var managedSession = new ManagedSession(options);
43+
44+
var session1 = new Session { Name = "session1", CreateTime = clock.GetCurrentDateTimeUtc().ToTimestamp() };
45+
client.Configure()
46+
.CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>())
47+
.Returns(Task.FromResult(session1));
48+
49+
await managedSession.EnsureFreshAsync(CancellationToken.None);
50+
51+
await client.Received(1).CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>());
52+
}
53+
54+
[Fact]
55+
public async Task GetFreshSession_NoExpiry()
56+
{
57+
var clock = new FakeClock(new DateTime(2025, 1, 1, 0, 0, 0, DateTimeKind.Utc));
58+
var client = SpannerClientHelpers.CreateMockClient(Logger.DefaultLogger, clock);
59+
60+
var options = ManagedSessionOptions.Create(SampleDatabaseName, client);
61+
var managedSession = new ManagedSession(options);
62+
63+
var session1 = new Session { Name = "session1", CreateTime = clock.GetCurrentDateTimeUtc().ToTimestamp() };
64+
client.Configure()
65+
.CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>())
66+
.Returns(Task.FromResult(session1));
67+
68+
await managedSession.EnsureFreshAsync(CancellationToken.None);
69+
70+
// Advance clock a bit, but not enough for soft expiry
71+
clock.Advance(TimeSpan.FromDays(1));
72+
await managedSession.EnsureFreshAsync(CancellationToken.None);
73+
74+
// Still only 1 call to CreateSessionAsync
75+
await client.Received(1).CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>());
76+
}
77+
78+
[Fact]
79+
public async Task GetFreshSession_SoftExpiry_BackgroundRefresh()
80+
{
81+
var clock = new FakeClock(new DateTime(2025, 1, 1, 0, 0, 0, DateTimeKind.Utc));
82+
var client = SpannerClientHelpers.CreateMockClient(Logger.DefaultLogger, clock);
83+
84+
var options = ManagedSessionOptions.Create(SampleDatabaseName, client);
85+
var managedSession = new ManagedSession(options);
86+
87+
var session1 = new Session { Name = "session1", CreateTime = clock.GetCurrentDateTimeUtc().ToTimestamp() };
88+
var session2 = new Session { Name = "session2" }; // CreateTime will be set when it's created
89+
90+
var refreshStartedTsc = new TaskCompletionSource<bool>();
91+
var freshSessionTsc = new TaskCompletionSource<Session>();
92+
client.Configure()
93+
.CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>())
94+
.Returns(Task.FromResult(session1), CreateSessionMock(refreshStartedTsc, freshSessionTsc.Task));
95+
96+
await managedSession.EnsureFreshAsync(CancellationToken.None);
97+
98+
// Advance past soft expiry
99+
clock.Advance(SoftExpiry + TimeSpan.FromMinutes(1));
100+
session2.CreateTime = clock.GetCurrentDateTimeUtc().ToTimestamp();
101+
102+
// This should trigger background refresh but return immediately because session1 is only soft expired.
103+
await managedSession.EnsureFreshAsync(CancellationToken.None);
104+
105+
// Wait for the background task to be triggered
106+
await refreshStartedTsc.Task;
107+
108+
await client.Received(2).CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>());
109+
110+
// Make the mock return so that the background refresh can complete.
111+
freshSessionTsc.SetResult(session2);
112+
113+
// Wait a little in real time to allow the background refresh to complete.
114+
await Task.Delay(100);
115+
116+
// Next call should return session2 (indirectly verified by no more calls to CreateSessionAsync)
117+
await managedSession.EnsureFreshAsync(CancellationToken.None);
118+
await client.Received(2).CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>());
119+
}
120+
121+
[Fact]
122+
public async Task GetFreshSession_ConcurrentCalls_OnlyOneRefresh()
123+
{
124+
var clock = new FakeClock(new DateTime(2025, 1, 1, 0, 0, 0, DateTimeKind.Utc));
125+
var client = SpannerClientHelpers.CreateMockClient(Logger.DefaultLogger, clock);
126+
127+
var options = ManagedSessionOptions.Create(SampleDatabaseName, client);
128+
var managedSession = new ManagedSession(options);
129+
130+
var session1 = new Session { Name = "session1", CreateTime = clock.GetCurrentDateTimeUtc().ToTimestamp() };
131+
var session2 = new Session { Name = "session2" };
132+
133+
var freshSessionTsc = new TaskCompletionSource<Session>();
134+
client.Configure().CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>())
135+
.Returns(Task.FromResult(session1), freshSessionTsc.Task);
136+
137+
await managedSession.EnsureFreshAsync(CancellationToken.None);
138+
139+
// Advance past hard expiry
140+
clock.Advance(HardExpiry + TimeSpan.FromMinutes(1));
141+
142+
// Start multiple concurrent calls
143+
var task1 = managedSession.EnsureFreshAsync(CancellationToken.None);
144+
var task2 = managedSession.EnsureFreshAsync(CancellationToken.None);
145+
var task3 = managedSession.EnsureFreshAsync(CancellationToken.None);
146+
147+
Assert.False(task1.IsCompleted);
148+
Assert.False(task2.IsCompleted);
149+
Assert.False(task3.IsCompleted);
150+
151+
freshSessionTsc.SetResult(session2);
152+
await Task.WhenAll(task1, task2, task3);
153+
154+
// Only 2 calls total (1 initial, 1 refresh)
155+
await client.Received(2).CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>());
156+
}
157+
158+
[Fact]
159+
public async Task GetFreshSession_SoftExpiry_RefreshFailure_Swallowed()
160+
{
161+
var clock = new FakeClock(new DateTime(2025, 1, 1, 0, 0, 0, DateTimeKind.Utc));
162+
var client = SpannerClientHelpers.CreateMockClient(Logger.DefaultLogger, clock);
163+
164+
var options = ManagedSessionOptions.Create(SampleDatabaseName, client);
165+
var managedSession = new ManagedSession(options);
166+
167+
var session1 = new Session { Name = "session1", CreateTime = clock.GetCurrentDateTimeUtc().ToTimestamp() };
168+
var session2 = new Session { Name = "session2" };
169+
170+
var firstRefreshStartedTcs = new TaskCompletionSource<bool>();
171+
var firstFreshSessionTcs = new TaskCompletionSource<Session>();
172+
var secondRefreshStartedTcs = new TaskCompletionSource<bool>();
173+
client.Configure()
174+
.CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>())
175+
.Returns(
176+
Task.FromResult(session1),
177+
CreateSessionMock(firstRefreshStartedTcs, firstFreshSessionTcs.Task),
178+
CreateSessionMock(secondRefreshStartedTcs, Task.FromResult(session2)));
179+
180+
await managedSession.EnsureFreshAsync(CancellationToken.None);
181+
182+
// Advance past soft expiry
183+
clock.Advance(SoftExpiry + TimeSpan.FromMinutes(1));
184+
session2.CreateTime = clock.GetCurrentDateTimeUtc().ToTimestamp();
185+
186+
// Should return session1 and not throw
187+
await managedSession.EnsureFreshAsync(CancellationToken.None);
188+
189+
// Wait for the background task to be triggered
190+
await firstRefreshStartedTcs.Task;
191+
192+
await client.Received(2).CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>());
193+
194+
// Complete with failure
195+
firstFreshSessionTcs.SetException(new Exception("Soft failure"));
196+
197+
// Wait a little in real time to allow the failed background refresh to complete.
198+
await Task.Delay(100);
199+
200+
// Still should not throw.
201+
// Note that since the previous background refresh failed, this call will trigger a new background refresh
202+
// because the session is still soft-expired.
203+
await managedSession.EnsureFreshAsync(CancellationToken.None);
204+
205+
// Wait for the background task to be triggered
206+
await secondRefreshStartedTcs.Task;
207+
await client.Received(3).CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>());
208+
}
209+
210+
[Fact]
211+
public async Task GetFreshSession_HardExpiry_BlockingRefresh()
212+
{
213+
var clock = new FakeClock(new DateTime(2025, 1, 1, 0, 0, 0, DateTimeKind.Utc));
214+
var client = SpannerClientHelpers.CreateMockClient(Logger.DefaultLogger, clock);
215+
216+
var options = ManagedSessionOptions.Create(SampleDatabaseName, client);
217+
var managedSession = new ManagedSession(options);
218+
219+
var session1 = new Session { Name = "session1", CreateTime = clock.GetCurrentDateTimeUtc().ToTimestamp() };
220+
var session2 = new Session { Name = "session2" };
221+
222+
var freshSessionTsc = new TaskCompletionSource<Session>();
223+
client.Configure().CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>())
224+
.Returns(Task.FromResult(session1), freshSessionTsc.Task);
225+
226+
await managedSession.EnsureFreshAsync(CancellationToken.None);
227+
228+
// Advance past hard expiry
229+
clock.Advance(HardExpiry + TimeSpan.FromMinutes(1));
230+
session2.CreateTime = clock.GetCurrentDateTimeUtc().ToTimestamp();
231+
232+
// This should block
233+
var refreshTask = managedSession.EnsureFreshAsync(CancellationToken.None);
234+
Assert.False(refreshTask.IsCompleted);
235+
236+
freshSessionTsc.SetResult(session2);
237+
await refreshTask;
238+
239+
await client.Received(2).CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>());
240+
}
241+
242+
[Fact]
243+
public async Task GetFreshSession_HardExpiry_RefreshFailure_Propagated()
244+
{
245+
var clock = new FakeClock(new DateTime(2025, 1, 1, 0, 0, 0, DateTimeKind.Utc));
246+
var client = SpannerClientHelpers.CreateMockClient(Logger.DefaultLogger, clock);
247+
248+
var options = ManagedSessionOptions.Create(SampleDatabaseName, client);
249+
var managedSession = new ManagedSession(options);
250+
251+
var session1 = new Session { Name = "session1", CreateTime = clock.GetCurrentDateTimeUtc().ToTimestamp() };
252+
253+
client.Configure().CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>())
254+
.Returns(Task.FromResult(session1), Task.FromException<Session>(new Exception("Hard failure")));
255+
256+
await managedSession.EnsureFreshAsync(CancellationToken.None);
257+
258+
// Advance past hard expiry
259+
clock.Advance(HardExpiry + TimeSpan.FromMinutes(1));
260+
261+
// Should throw
262+
var ex = await Assert.ThrowsAsync<Exception>(() => managedSession.EnsureFreshAsync(CancellationToken.None));
263+
Assert.Equal("Hard failure", ex.Message);
264+
}
265+
266+
[Fact]
267+
public async Task BeginTransactionAsync_ConcurrentCalls()
268+
{
269+
var clock = new FakeClock(new DateTime(2025, 1, 1, 0, 0, 0, DateTimeKind.Utc));
270+
var client = SpannerClientHelpers.CreateMockClient(Logger.DefaultLogger, clock);
271+
272+
var options = ManagedSessionOptions.Create(SampleDatabaseName, client);
273+
var managedSession = new ManagedSession(options);
274+
275+
var session1 = new Session { Name = "session1", CreateTime = clock.GetCurrentDateTimeUtc().ToTimestamp() };
276+
var session2 = new Session { Name = "session2" };
277+
278+
var tcs = new TaskCompletionSource<Session>();
279+
client.Configure().CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>())
280+
.Returns(Task.FromResult(session1), tcs.Task);
281+
282+
// Initial creation
283+
await managedSession.EnsureFreshAsync(CancellationToken.None);
284+
285+
// Advance past hard expiry
286+
clock.Advance(HardExpiry + TimeSpan.FromMinutes(1));
287+
288+
// Start multiple concurrent calls to BeginTransactionAsync
289+
var task1 = managedSession.BeginTransactionAsync(new TransactionOptions(), false, false, CancellationToken.None);
290+
var task2 = managedSession.BeginTransactionAsync(new TransactionOptions(), false, false, CancellationToken.None);
291+
var task3 = managedSession.BeginTransactionAsync(new TransactionOptions(), false, false, CancellationToken.None);
292+
293+
Assert.False(task1.IsCompleted);
294+
Assert.False(task2.IsCompleted);
295+
Assert.False(task3.IsCompleted);
296+
297+
tcs.SetResult(session2);
298+
var transactions = await Task.WhenAll(task1, task2, task3);
299+
300+
// Only 2 calls total (1 initial, 1 refresh)
301+
await client.Received(2).CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>());
302+
303+
// Verify all transactions use the new session
304+
foreach (var transaction in transactions)
305+
{
306+
Assert.Equal(session2.SessionName, transaction.SessionName);
307+
}
308+
}
309+
310+
[Fact]
311+
public async Task BatchWriteAsync_EnsureFreshnessAndSetsSession()
312+
{
313+
var clock = new FakeClock(new DateTime(2025, 1, 1, 0, 0, 0, DateTimeKind.Utc));
314+
var client = SpannerClientHelpers.CreateMockClient(Logger.DefaultLogger, clock);
315+
316+
var options = ManagedSessionOptions.Create(SampleDatabaseName, client);
317+
var managedSession = new ManagedSession(options);
318+
319+
var session1 = new Session { Name = "session1", CreateTime = clock.GetCurrentDateTimeUtc().ToTimestamp() };
320+
client.Configure().CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>())
321+
.Returns(Task.FromResult(session1));
322+
323+
// Mock BatchWrite
324+
var batchWriteStream = Substitute.For<SpannerClient.BatchWriteStream>();
325+
client.BatchWrite(Arg.Any<BatchWriteRequest>(), Arg.Any<CallSettings>())
326+
.Returns(batchWriteStream);
327+
328+
// Execute BatchWriteAsync
329+
var request = new BatchWriteRequest();
330+
await managedSession.BatchWriteAsync(request, null);
331+
332+
// Verify session was created
333+
await client.Received(1).CreateSessionAsync(Arg.Any<CreateSessionRequest>(), Arg.Any<CallSettings>());
334+
335+
// Verify BatchWrite was called with the correct session
336+
client.Received(1).BatchWrite(Arg.Is<BatchWriteRequest>(r => r.Session == session1.Name), Arg.Any<CallSettings>());
337+
}
338+
339+
private static Task<Session> CreateSessionMock(TaskCompletionSource<bool> methodStartedTsc, Task<Session> result) =>
340+
Task.Run(() =>
341+
{
342+
methodStartedTsc.SetResult(true);
343+
return result;
344+
});
345+
}

0 commit comments

Comments
 (0)