Skip to content

Commit 1ca0250

Browse files
ivanauthivanauth
authored
test: use synctest in some CRDB pool retry tests (#2754)
1 parent f0f94ee commit 1ca0250

File tree

1 file changed

+91
-73
lines changed

1 file changed

+91
-73
lines changed

‎internal/datastore/crdb/pool/pool_test.go‎

Lines changed: 91 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"context"
55
"errors"
66
"testing"
7+
"testing/synctest"
78
"time"
89

910
"github.com/jackc/pgx/v5"
@@ -79,58 +80,65 @@ func createTestRetryPool(testPool *TestPool) *RetryPool {
7980
}
8081

8182
func TestContextCancelledDuringBlockingAcquire(t *testing.T) {
82-
testPool := NewTestPool()
83-
testPool.acquireFunc = func(ctx context.Context) (*pgxpool.Conn, error) {
84-
// Block until context is cancelled
85-
<-ctx.Done()
86-
return nil, ctx.Err()
87-
}
83+
t.Parallel()
84+
synctest.Test(t, func(t *testing.T) {
85+
testPool := NewTestPool()
86+
testPool.acquireFunc = func(ctx context.Context) (*pgxpool.Conn, error) {
87+
// Block until context is cancelled
88+
<-ctx.Done()
89+
return nil, ctx.Err()
90+
}
8891

89-
retryPool := createTestRetryPool(testPool)
90-
ctx, cancel := context.WithCancel(context.Background())
92+
retryPool := createTestRetryPool(testPool)
93+
ctx, cancel := context.WithCancel(t.Context())
9194

92-
// Cancel the context after a short delay
93-
go func() {
94-
time.Sleep(10 * time.Millisecond)
95-
cancel()
96-
}()
95+
// Cancel the context after a short delay
96+
go func() {
97+
time.Sleep(10 * time.Millisecond)
98+
cancel()
99+
}()
97100

98-
err := retryPool.withRetries(ctx, 0, func(conn *pgxpool.Conn) error {
99-
t.Fatal("function should not be called when acquire fails")
100-
return nil
101-
})
101+
err := retryPool.withRetries(ctx, 0, func(conn *pgxpool.Conn) error {
102+
t.Fatal("function should not be called when acquire fails")
103+
return nil
104+
})
102105

103-
require.Error(t, err)
104-
require.ErrorIs(t, err, context.Canceled)
106+
require.Error(t, err)
107+
require.ErrorIs(t, err, context.Canceled)
108+
})
105109
}
106110

107111
func TestAcquireTimeoutReturnsErrAcquire(t *testing.T) {
108-
testPool := NewTestPool()
109-
testPool.acquireFunc = func(ctx context.Context) (*pgxpool.Conn, error) {
110-
// Simulate slow acquire that times out
111-
select {
112-
case <-time.After(100 * time.Millisecond):
113-
return nil, errors.New("should not reach here")
114-
case <-ctx.Done():
115-
return nil, ctx.Err()
112+
t.Parallel()
113+
synctest.Test(t, func(t *testing.T) {
114+
testPool := NewTestPool()
115+
testPool.acquireFunc = func(ctx context.Context) (*pgxpool.Conn, error) {
116+
// Simulate slow acquire that times out
117+
select {
118+
case <-time.After(100 * time.Millisecond):
119+
return nil, errors.New("should not reach here")
120+
case <-ctx.Done():
121+
return nil, ctx.Err()
122+
}
116123
}
117-
}
118124

119-
retryPool := createTestRetryPool(testPool)
120-
ctx := context.Background()
121-
acquireTimeout := 50 * time.Millisecond
125+
retryPool := createTestRetryPool(testPool)
126+
ctx := t.Context()
127+
acquireTimeout := 50 * time.Millisecond
122128

123-
err := retryPool.withRetries(ctx, acquireTimeout, func(conn *pgxpool.Conn) error {
124-
t.Fatal("function should not be called when acquire times out")
125-
return nil
126-
})
129+
err := retryPool.withRetries(ctx, acquireTimeout, func(conn *pgxpool.Conn) error {
130+
t.Fatal("function should not be called when acquire times out")
131+
return nil
132+
})
127133

128-
require.Error(t, err)
129-
require.Contains(t, err.Error(), "error acquiring connection from pool")
130-
require.ErrorIs(t, errors.Unwrap(err), ErrAcquire)
134+
require.Error(t, err)
135+
require.Contains(t, err.Error(), "error acquiring connection from pool")
136+
require.ErrorIs(t, errors.Unwrap(err), ErrAcquire)
137+
})
131138
}
132139

133140
func TestAcquireSucceedsButTopLevelContextCancelled(t *testing.T) {
141+
t.Parallel()
134142
testPool := NewTestPool()
135143

136144
retryPool := createTestRetryPool(testPool)
@@ -147,6 +155,7 @@ func TestAcquireSucceedsButTopLevelContextCancelled(t *testing.T) {
147155
}
148156

149157
func TestAcquireErrorWithConnectionReturned(t *testing.T) {
158+
t.Parallel()
150159
testPool := NewTestPool()
151160
testPool.acquireFunc = func(ctx context.Context) (*pgxpool.Conn, error) {
152161
// Return both connection and error
@@ -167,32 +176,36 @@ func TestAcquireErrorWithConnectionReturned(t *testing.T) {
167176
}
168177

169178
func TestAcquireSucceedsWithinTimeout(t *testing.T) {
170-
testPool := NewTestPool()
171-
testPool.acquireFunc = func(ctx context.Context) (*pgxpool.Conn, error) {
172-
// Small delay but within timeout
173-
select {
174-
case <-time.After(10 * time.Millisecond):
175-
return &pgxpool.Conn{}, nil
176-
case <-ctx.Done():
177-
return nil, ctx.Err()
179+
t.Parallel()
180+
synctest.Test(t, func(t *testing.T) {
181+
testPool := NewTestPool()
182+
testPool.acquireFunc = func(ctx context.Context) (*pgxpool.Conn, error) {
183+
// Small delay but within timeout
184+
select {
185+
case <-time.After(10 * time.Millisecond):
186+
return &pgxpool.Conn{}, nil
187+
case <-ctx.Done():
188+
return nil, ctx.Err()
189+
}
178190
}
179-
}
180191

181-
retryPool := createTestRetryPool(testPool)
182-
ctx := context.Background()
183-
acquireTimeout := 50 * time.Millisecond
184-
functionCalled := false
192+
retryPool := createTestRetryPool(testPool)
193+
ctx := t.Context()
194+
acquireTimeout := 50 * time.Millisecond
195+
functionCalled := false
185196

186-
err := retryPool.withRetries(ctx, acquireTimeout, func(conn *pgxpool.Conn) error {
187-
functionCalled = true
188-
return nil
189-
})
197+
err := retryPool.withRetries(ctx, acquireTimeout, func(conn *pgxpool.Conn) error {
198+
functionCalled = true
199+
return nil
200+
})
190201

191-
require.NoError(t, err)
192-
require.True(t, functionCalled, "function should have been called")
202+
require.NoError(t, err)
203+
require.True(t, functionCalled, "function should have been called")
204+
})
193205
}
194206

195207
func TestNoAcquireTimeoutUsesOriginalContext(t *testing.T) {
208+
t.Parallel()
196209
var acquireContext context.Context
197210

198211
testPool := NewTestPool()
@@ -213,6 +226,7 @@ func TestNoAcquireTimeoutUsesOriginalContext(t *testing.T) {
213226
}
214227

215228
func TestAcquireTimeoutCreatesSeparateContext(t *testing.T) {
229+
t.Parallel()
216230
var acquireContext context.Context
217231

218232
testPool := NewTestPool()
@@ -241,28 +255,32 @@ func TestAcquireTimeoutCreatesSeparateContext(t *testing.T) {
241255
}
242256

243257
func TestAcquireTimeoutContextCausePreserved(t *testing.T) {
244-
testPool := NewTestPool()
245-
testPool.acquireFunc = func(ctx context.Context) (*pgxpool.Conn, error) {
246-
// Wait for context timeout
247-
<-ctx.Done()
248-
return nil, ctx.Err()
249-
}
258+
t.Parallel()
259+
synctest.Test(t, func(t *testing.T) {
260+
testPool := NewTestPool()
261+
testPool.acquireFunc = func(ctx context.Context) (*pgxpool.Conn, error) {
262+
// Wait for context timeout
263+
<-ctx.Done()
264+
return nil, ctx.Err()
265+
}
250266

251-
retryPool := createTestRetryPool(testPool)
252-
ctx := context.Background()
253-
acquireTimeout := 10 * time.Millisecond
267+
retryPool := createTestRetryPool(testPool)
268+
ctx := t.Context()
269+
acquireTimeout := 10 * time.Millisecond
254270

255-
err := retryPool.withRetries(ctx, acquireTimeout, func(conn *pgxpool.Conn) error {
256-
t.Fatal("function should not be called")
257-
return nil
258-
})
271+
err := retryPool.withRetries(ctx, acquireTimeout, func(conn *pgxpool.Conn) error {
272+
t.Fatal("function should not be called")
273+
return nil
274+
})
259275

260-
require.Error(t, err)
261-
require.Contains(t, err.Error(), "error acquiring connection from pool")
262-
require.ErrorIs(t, errors.Unwrap(err), ErrAcquire)
276+
require.Error(t, err)
277+
require.Contains(t, err.Error(), "error acquiring connection from pool")
278+
require.ErrorIs(t, errors.Unwrap(err), ErrAcquire)
279+
})
263280
}
264281

265282
func TestSuccessfulFunctionExecution(t *testing.T) {
283+
t.Parallel()
266284
testPool := NewTestPool()
267285
testPool.acquireFunc = func(ctx context.Context) (*pgxpool.Conn, error) {
268286
return &pgxpool.Conn{}, nil

0 commit comments

Comments
 (0)