Skip to content

Commit b7d70ac

Browse files
dgageotdgageot
committed
Use t.Context()
Signed-off-by: David Gageot <david.gageot@docker.com>
1 parent d81d691 commit b7d70ac

9 files changed

Lines changed: 61 additions & 67 deletions

File tree

‎pkg/memory/database/sqlite/sqlite_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ func TestAddMemory(t *testing.T) {
4646
db, _, cleanup := setupTestDB(t)
4747
defer cleanup()
4848

49-
ctx := context.Background()
49+
ctx := t.Context()
5050

5151
// Test adding a valid memory
5252
memory := database.UserMemory{
@@ -165,7 +165,7 @@ func TestDatabaseOperationsWithCanceledContext(t *testing.T) {
165165
defer cleanup()
166166

167167
// Create a canceled context
168-
ctx, cancel := context.WithCancel(context.Background())
168+
ctx, cancel := context.WithCancel(t.Context())
169169
cancel()
170170

171171
// Test operations with canceled context

‎pkg/servicecore/store_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ func TestSQLiteStore_Migration(t *testing.T) {
2121
defer store.(*SQLiteStore).Close()
2222

2323
// Verify migration was applied by creating a client
24-
ctx := context.Background()
24+
ctx := t.Context()
2525
err = store.CreateClient(ctx, "test-client")
2626
assert.NoError(t, err)
2727
}

‎pkg/session/store_test.go‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package session
22

33
import (
4-
"context"
54
"os"
65
"testing"
76
"time"
@@ -47,11 +46,11 @@ func TestStoreAgentName(t *testing.T) {
4746
}
4847

4948
// Store the session
50-
err = store.AddSession(context.Background(), session)
49+
err = store.AddSession(t.Context(), session)
5150
require.NoError(t, err)
5251

5352
// Retrieve the session
54-
retrievedSession, err := store.GetSession(context.Background(), "test-session")
53+
retrievedSession, err := store.GetSession(t.Context(), "test-session")
5554
require.NoError(t, err)
5655
require.NotNil(t, retrievedSession)
5756

@@ -106,11 +105,11 @@ func TestStoreMultipleAgents(t *testing.T) {
106105
}
107106

108107
// Store the session
109-
err = store.AddSession(context.Background(), session)
108+
err = store.AddSession(t.Context(), session)
110109
require.NoError(t, err)
111110

112111
// Retrieve the session
113-
retrievedSession, err := store.GetSession(context.Background(), "multi-agent-session")
112+
retrievedSession, err := store.GetSession(t.Context(), "multi-agent-session")
114113
require.NoError(t, err)
115114
require.NotNil(t, retrievedSession)
116115

@@ -169,13 +168,13 @@ func TestGetSessions(t *testing.T) {
169168
}
170169

171170
// Store the sessions
172-
err = store.AddSession(context.Background(), session1)
171+
err = store.AddSession(t.Context(), session1)
173172
require.NoError(t, err)
174-
err = store.AddSession(context.Background(), session2)
173+
err = store.AddSession(t.Context(), session2)
175174
require.NoError(t, err)
176175

177176
// Retrieve all sessions
178-
sessions, err := store.GetSessions(context.Background())
177+
sessions, err := store.GetSessions(t.Context())
179178
require.NoError(t, err)
180179
assert.Len(t, sessions, 2)
181180

@@ -218,11 +217,11 @@ func TestStoreAgentNameJSON(t *testing.T) {
218217
}
219218

220219
// Store the session
221-
err = store.AddSession(context.Background(), session)
220+
err = store.AddSession(t.Context(), session)
222221
require.NoError(t, err)
223222

224223
// Retrieve the session
225-
retrievedSession, err := store.GetSession(context.Background(), "json-test-session")
224+
retrievedSession, err := store.GetSession(t.Context(), "json-test-session")
226225
require.NoError(t, err)
227226
require.NotNil(t, retrievedSession)
228227

‎pkg/tools/builtin/filesystem_test.go‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package builtin
22

33
import (
4-
"context"
54
"encoding/json"
65
"os"
76
"path/filepath"
@@ -33,7 +32,7 @@ func TestFilesystemTool_Instructions(t *testing.T) {
3332

3433
func TestFilesystemTool_Tools(t *testing.T) {
3534
tool := NewFilesystemTool([]string{"/tmp"})
36-
tools, err := tool.Tools(context.Background())
35+
tools, err := tool.Tools(t.Context())
3736

3837
require.NoError(t, err)
3938
assert.Len(t, tools, 14)
@@ -570,7 +569,7 @@ func TestFilesystemTool_InvalidArguments(t *testing.T) {
570569
},
571570
}
572571

573-
result, err := handler(context.Background(), toolCall)
572+
result, err := handler(t.Context(), toolCall)
574573
assert.Error(t, err)
575574
assert.Nil(t, result)
576575
}
@@ -579,7 +578,7 @@ func TestFilesystemTool_StartStop(t *testing.T) {
579578
tool := NewFilesystemTool([]string{"/tmp"})
580579

581580
// Test Start method
582-
err := tool.Start(context.Background())
581+
err := tool.Start(t.Context())
583582
require.NoError(t, err)
584583

585584
// Test Stop method
@@ -590,7 +589,7 @@ func TestFilesystemTool_StartStop(t *testing.T) {
590589
// Helper functions
591590

592591
func getToolHandler(t *testing.T, tool *FilesystemTool, toolName string) tools.ToolHandler {
593-
tls, err := tool.Tools(context.Background())
592+
tls, err := tool.Tools(t.Context())
594593
require.NoError(t, err)
595594

596595
for _, tl := range tls {
@@ -613,7 +612,7 @@ func callHandler(t *testing.T, handler tools.ToolHandler, args map[string]any) *
613612
},
614613
}
615614

616-
result, err := handler(context.Background(), toolCall)
615+
result, err := handler(t.Context(), toolCall)
617616
require.NoError(t, err)
618617
require.NotNil(t, result)
619618

‎pkg/tools/builtin/memory_test.go‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ func TestMemoryTool_Tools(t *testing.T) {
5454
manager := new(MockMemoryManager)
5555
tool := NewMemoryTool(manager)
5656

57-
tls, err := tool.Tools(context.Background())
57+
tls, err := tool.Tools(t.Context())
5858

5959
require.NoError(t, err)
6060
assert.Len(t, tls, 3)
@@ -90,7 +90,7 @@ func TestMemoryTool_HandleAddMemory(t *testing.T) {
9090
}
9191

9292
// Call handler
93-
result, err := tool.handleAddMemory(context.Background(), toolCall)
93+
result, err := tool.handleAddMemory(t.Context(), toolCall)
9494

9595
// Verify
9696
require.NoError(t, err)
@@ -126,7 +126,7 @@ func TestMemoryTool_HandleGetMemories(t *testing.T) {
126126
}
127127

128128
// Call handler
129-
result, err := tool.handleGetMemories(context.Background(), toolCall)
129+
result, err := tool.handleGetMemories(t.Context(), toolCall)
130130

131131
// Verify
132132
require.NoError(t, err)
@@ -165,7 +165,7 @@ func TestMemoryTool_HandleDeleteMemory(t *testing.T) {
165165
}
166166

167167
// Call handler
168-
result, err := tool.handleDeleteMemory(context.Background(), toolCall)
168+
result, err := tool.handleDeleteMemory(t.Context(), toolCall)
169169

170170
// Verify
171171
require.NoError(t, err)
@@ -185,7 +185,7 @@ func TestMemoryTool_InvalidArguments(t *testing.T) {
185185
},
186186
}
187187

188-
result, err := tool.handleAddMemory(context.Background(), toolCall)
188+
result, err := tool.handleAddMemory(t.Context(), toolCall)
189189
assert.Error(t, err)
190190
assert.Nil(t, result)
191191

@@ -197,7 +197,7 @@ func TestMemoryTool_InvalidArguments(t *testing.T) {
197197
},
198198
}
199199

200-
result, err = tool.handleDeleteMemory(context.Background(), toolCall)
200+
result, err = tool.handleDeleteMemory(t.Context(), toolCall)
201201
assert.Error(t, err)
202202
assert.Nil(t, result)
203203
}
@@ -207,7 +207,7 @@ func TestMemoryTool_StartStop(t *testing.T) {
207207
tool := NewMemoryTool(manager)
208208

209209
// Test Start method
210-
err := tool.Start(context.Background())
210+
err := tool.Start(t.Context())
211211
require.NoError(t, err)
212212

213213
// Test Stop method

‎pkg/tools/builtin/shell_test.go‎

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package builtin
22

33
import (
4-
"context"
54
"encoding/json"
65
"os"
76
"strings"
@@ -38,7 +37,7 @@ func TestNewShellTool(t *testing.T) {
3837
func TestShellTool_Tools(t *testing.T) {
3938
tool := NewShellTool()
4039

41-
tools, err := tool.Tools(context.Background())
40+
tools, err := tool.Tools(t.Context())
4241

4342
require.NoError(t, err)
4443
assert.Len(t, tools, 1)
@@ -65,7 +64,7 @@ func TestShellTool_HandlerEcho(t *testing.T) {
6564
tool := NewShellTool()
6665

6766
// Get handler from tool
68-
tls, err := tool.Tools(context.Background())
67+
tls, err := tool.Tools(t.Context())
6968
require.NoError(t, err)
7069
require.Len(t, tls, 1)
7170

@@ -89,7 +88,7 @@ func TestShellTool_HandlerEcho(t *testing.T) {
8988
}
9089

9190
// Call handler
92-
result, err := handler(context.Background(), toolCall)
91+
result, err := handler(t.Context(), toolCall)
9392

9493
// Verify
9594
require.NoError(t, err)
@@ -101,7 +100,7 @@ func TestShellTool_HandlerWithCwd(t *testing.T) {
101100
tool := NewShellTool()
102101

103102
// Get handler from tool
104-
tls, err := tool.Tools(context.Background())
103+
tls, err := tool.Tools(t.Context())
105104
require.NoError(t, err)
106105
require.Len(t, tls, 1)
107106

@@ -127,7 +126,7 @@ func TestShellTool_HandlerWithCwd(t *testing.T) {
127126
}
128127

129128
// Call handler
130-
result, err := handler(context.Background(), toolCall)
129+
result, err := handler(t.Context(), toolCall)
131130

132131
// Verify
133132
require.NoError(t, err)
@@ -142,7 +141,7 @@ func TestShellTool_HandlerError(t *testing.T) {
142141
tool := NewShellTool()
143142

144143
// Get handler from tool
145-
tls, err := tool.Tools(context.Background())
144+
tls, err := tool.Tools(t.Context())
146145
require.NoError(t, err)
147146
require.Len(t, tls, 1)
148147

@@ -166,7 +165,7 @@ func TestShellTool_HandlerError(t *testing.T) {
166165
}
167166

168167
// Call handler
169-
result, err := handler(context.Background(), toolCall)
168+
result, err := handler(t.Context(), toolCall)
170169

171170
// Verify
172171
require.NoError(t, err, "Handler should not return an error")
@@ -177,7 +176,7 @@ func TestShellTool_InvalidArguments(t *testing.T) {
177176
tool := NewShellTool()
178177

179178
// Get handler from tool
180-
tls, err := tool.Tools(context.Background())
179+
tls, err := tool.Tools(t.Context())
181180
require.NoError(t, err)
182181
require.Len(t, tls, 1)
183182

@@ -191,7 +190,7 @@ func TestShellTool_InvalidArguments(t *testing.T) {
191190
},
192191
}
193192

194-
result, err := handler(context.Background(), toolCall)
193+
result, err := handler(t.Context(), toolCall)
195194
assert.Error(t, err)
196195
assert.Nil(t, result)
197196
}
@@ -200,7 +199,7 @@ func TestShellTool_StartStop(t *testing.T) {
200199
tool := NewShellTool()
201200

202201
// Test Start method
203-
err := tool.Start(context.Background())
202+
err := tool.Start(t.Context())
204203
require.NoError(t, err)
205204

206205
// Test Stop method

‎pkg/tools/builtin/think_test.go‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package builtin
22

33
import (
4-
"context"
54
"encoding/json"
65
"testing"
76

@@ -30,7 +29,7 @@ func TestThinkTool_Instructions(t *testing.T) {
3029
func TestThinkTool_Tools(t *testing.T) {
3130
tool := NewThinkTool()
3231

33-
tls, err := tool.Tools(context.Background())
32+
tls, err := tool.Tools(t.Context())
3433

3534
require.NoError(t, err)
3635
assert.Len(t, tls, 1)
@@ -54,7 +53,7 @@ func TestThinkTool_Handler(t *testing.T) {
5453
tool := NewThinkTool()
5554

5655
// Get handler from tool
57-
tls, err := tool.Tools(context.Background())
56+
tls, err := tool.Tools(t.Context())
5857
require.NoError(t, err)
5958
require.Len(t, tls, 1)
6059

@@ -76,7 +75,7 @@ func TestThinkTool_Handler(t *testing.T) {
7675
}
7776

7877
// Call handler
79-
result, err := handler(context.Background(), toolCall)
78+
result, err := handler(t.Context(), toolCall)
8079

8180
// Verify
8281
require.NoError(t, err)
@@ -88,7 +87,7 @@ func TestThinkTool_Handler(t *testing.T) {
8887

8988
toolCall.Function.Arguments = string(argsBytes)
9089

91-
result, err = handler(context.Background(), toolCall)
90+
result, err = handler(t.Context(), toolCall)
9291

9392
// Verify both thoughts are in output
9493
require.NoError(t, err)
@@ -100,7 +99,7 @@ func TestThinkTool_InvalidArguments(t *testing.T) {
10099
tool := NewThinkTool()
101100

102101
// Get handler from tool
103-
tls, err := tool.Tools(context.Background())
102+
tls, err := tool.Tools(t.Context())
104103
require.NoError(t, err)
105104
require.Len(t, tls, 1)
106105

@@ -114,7 +113,7 @@ func TestThinkTool_InvalidArguments(t *testing.T) {
114113
},
115114
}
116115

117-
result, err := handler(context.Background(), toolCall)
116+
result, err := handler(t.Context(), toolCall)
118117
assert.Error(t, err)
119118
assert.Nil(t, result)
120119
}
@@ -123,7 +122,7 @@ func TestThinkTool_StartStop(t *testing.T) {
123122
tool := NewThinkTool()
124123

125124
// Test Start method
126-
err := tool.Start(context.Background())
125+
err := tool.Start(t.Context())
127126
require.NoError(t, err)
128127

129128
// Test Stop method

0 commit comments

Comments
 (0)