@@ -7,15 +7,20 @@ import (
7
7
"net"
8
8
"sync"
9
9
"testing"
10
+ "time"
10
11
11
12
"github.com/google/uuid"
12
13
"github.com/pion/udp"
13
14
"github.com/stretchr/testify/assert"
14
15
"github.com/stretchr/testify/require"
15
16
17
+ "cdr.dev/slog"
18
+ "cdr.dev/slog/sloggers/slogtest"
19
+ "github.com/coder/coder/v2/agent"
16
20
"github.com/coder/coder/v2/cli/clitest"
17
21
"github.com/coder/coder/v2/coderd/coderdtest"
18
22
"github.com/coder/coder/v2/codersdk"
23
+ "github.com/coder/coder/v2/codersdk/agentsdk"
19
24
"github.com/coder/coder/v2/provisioner/echo"
20
25
"github.com/coder/coder/v2/pty/ptytest"
21
26
"github.com/coder/coder/v2/testutil"
@@ -312,24 +317,23 @@ func runAgent(t *testing.T, client *codersdk.Client, userID uuid.UUID) codersdk.
312
317
workspace := coderdtest .CreateWorkspace (t , client , orgID , template .ID )
313
318
coderdtest .AwaitWorkspaceBuildJob (t , client , workspace .LatestBuild .ID )
314
319
315
- // Start workspace agent in a goroutine
316
- inv , root := clitest .New (t , "agent" , "--agent-token" , agentToken , "--agent-url" , client .URL .String ())
317
- clitest .SetupConfig (t , client , root )
318
- pty := ptytest .New (t )
319
- inv .Stdin = pty .Input ()
320
- inv .Stdout = pty .Output ()
321
- inv .Stderr = pty .Output ()
322
- errC := make (chan error )
323
- agentCtx , agentCancel := context .WithCancel (ctx )
320
+ logger := slogtest .Make (t , nil ).Leveled (slog .LevelDebug ).Named ("agent" )
321
+ agentClient := agentsdk .New (client .URL )
322
+ agentClient .SDK .SetLogger (logger )
323
+ agentClient .SDK .SetSessionToken (agentToken )
324
+ agnt := agent .New (agent.Options {
325
+ Client : agentClient ,
326
+ Logger : logger ,
327
+ LogDir : t .TempDir (),
328
+ ExchangeToken : func (ctx context.Context ) (string , error ) {
329
+ return agentToken , nil
330
+ },
331
+ SSHMaxTimeout : time .Second * 60 ,
332
+ })
324
333
t .Cleanup (func () {
325
- agentCancel ()
326
- err := <- errC
327
- require .NoError (t , err )
334
+ err := agnt .Close ()
335
+ assert .NoError (t , err )
328
336
})
329
- go func () {
330
- errC <- inv .WithContext (agentCtx ).Run ()
331
- }()
332
-
333
337
coderdtest .AwaitWorkspaceAgents (t , client , workspace .ID )
334
338
335
339
return workspace
0 commit comments