@@ -65,38 +65,48 @@ func TestMain(m *testing.M) {
65
65
goleak .VerifyTestMain (m , testutil .GoleakOptions ... )
66
66
}
67
67
68
+ var sshPorts = []uint16 {workspacesdk .AgentSSHPort , workspacesdk .AgentStandardSSHPort }
69
+
68
70
// NOTE: These tests only work when your default shell is bash for some reason.
69
71
70
72
func TestAgent_Stats_SSH (t * testing.T ) {
71
73
t .Parallel ()
72
- ctx , cancel := context .WithTimeout (context .Background (), testutil .WaitLong )
73
- defer cancel ()
74
74
75
- //nolint:dogsled
76
- conn , _ , stats , _ , _ := setupAgent (t , agentsdk.Manifest {}, 0 )
75
+ for _ , port := range sshPorts {
76
+ port := port
77
+ t .Run (fmt .Sprintf ("(:%d)" , port ), func (t * testing.T ) {
78
+ t .Parallel ()
77
79
78
- sshClient , err := conn .SSHClient (ctx )
79
- require .NoError (t , err )
80
- defer sshClient .Close ()
81
- session , err := sshClient .NewSession ()
82
- require .NoError (t , err )
83
- defer session .Close ()
84
- stdin , err := session .StdinPipe ()
85
- require .NoError (t , err )
86
- err = session .Shell ()
87
- require .NoError (t , err )
80
+ ctx , cancel := context .WithTimeout (context .Background (), testutil .WaitLong )
81
+ defer cancel ()
88
82
89
- var s * proto.Stats
90
- require .Eventuallyf (t , func () bool {
91
- var ok bool
92
- s , ok = <- stats
93
- return ok && s .ConnectionCount > 0 && s .RxBytes > 0 && s .TxBytes > 0 && s .SessionCountSsh == 1
94
- }, testutil .WaitLong , testutil .IntervalFast ,
95
- "never saw stats: %+v" , s ,
96
- )
97
- _ = stdin .Close ()
98
- err = session .Wait ()
99
- require .NoError (t , err )
83
+ //nolint:dogsled
84
+ conn , _ , stats , _ , _ := setupAgent (t , agentsdk.Manifest {}, 0 )
85
+
86
+ sshClient , err := conn .SSHClientOnPort (ctx , port )
87
+ require .NoError (t , err )
88
+ defer sshClient .Close ()
89
+ session , err := sshClient .NewSession ()
90
+ require .NoError (t , err )
91
+ defer session .Close ()
92
+ stdin , err := session .StdinPipe ()
93
+ require .NoError (t , err )
94
+ err = session .Shell ()
95
+ require .NoError (t , err )
96
+
97
+ var s * proto.Stats
98
+ require .Eventuallyf (t , func () bool {
99
+ var ok bool
100
+ s , ok = <- stats
101
+ return ok && s .ConnectionCount > 0 && s .RxBytes > 0 && s .TxBytes > 0 && s .SessionCountSsh == 1
102
+ }, testutil .WaitLong , testutil .IntervalFast ,
103
+ "never saw stats: %+v" , s ,
104
+ )
105
+ _ = stdin .Close ()
106
+ err = session .Wait ()
107
+ require .NoError (t , err )
108
+ })
109
+ }
100
110
}
101
111
102
112
func TestAgent_Stats_ReconnectingPTY (t * testing.T ) {
@@ -278,15 +288,23 @@ func TestAgent_Stats_Magic(t *testing.T) {
278
288
279
289
func TestAgent_SessionExec (t * testing.T ) {
280
290
t .Parallel ()
281
- session := setupSSHSession (t , agentsdk.Manifest {}, codersdk.ServiceBannerConfig {}, nil )
282
291
283
- command := "echo test"
284
- if runtime .GOOS == "windows" {
285
- command = "cmd.exe /c echo test"
292
+ for _ , port := range sshPorts {
293
+ port := port
294
+ t .Run (fmt .Sprintf ("(:%d)" , port ), func (t * testing.T ) {
295
+ t .Parallel ()
296
+
297
+ session := setupSSHSessionOnPort (t , agentsdk.Manifest {}, codersdk.ServiceBannerConfig {}, nil , port )
298
+
299
+ command := "echo test"
300
+ if runtime .GOOS == "windows" {
301
+ command = "cmd.exe /c echo test"
302
+ }
303
+ output , err := session .Output (command )
304
+ require .NoError (t , err )
305
+ require .Equal (t , "test" , strings .TrimSpace (string (output )))
306
+ })
286
307
}
287
- output , err := session .Output (command )
288
- require .NoError (t , err )
289
- require .Equal (t , "test" , strings .TrimSpace (string (output )))
290
308
}
291
309
292
310
//nolint:tparallel // Sub tests need to run sequentially.
@@ -396,25 +414,33 @@ func TestAgent_SessionTTYShell(t *testing.T) {
396
414
// it seems like it could be either.
397
415
t .Skip ("ConPTY appears to be inconsistent on Windows." )
398
416
}
399
- session := setupSSHSession (t , agentsdk.Manifest {}, codersdk.ServiceBannerConfig {}, nil )
400
- command := "sh"
401
- if runtime .GOOS == "windows" {
402
- command = "cmd.exe"
417
+
418
+ for _ , port := range sshPorts {
419
+ port := port
420
+ t .Run (fmt .Sprintf ("(%d)" , port ), func (t * testing.T ) {
421
+ t .Parallel ()
422
+
423
+ session := setupSSHSessionOnPort (t , agentsdk.Manifest {}, codersdk.ServiceBannerConfig {}, nil , port )
424
+ command := "sh"
425
+ if runtime .GOOS == "windows" {
426
+ command = "cmd.exe"
427
+ }
428
+ err := session .RequestPty ("xterm" , 128 , 128 , ssh.TerminalModes {})
429
+ require .NoError (t , err )
430
+ ptty := ptytest .New (t )
431
+ session .Stdout = ptty .Output ()
432
+ session .Stderr = ptty .Output ()
433
+ session .Stdin = ptty .Input ()
434
+ err = session .Start (command )
435
+ require .NoError (t , err )
436
+ _ = ptty .Peek (ctx , 1 ) // wait for the prompt
437
+ ptty .WriteLine ("echo test" )
438
+ ptty .ExpectMatch ("test" )
439
+ ptty .WriteLine ("exit" )
440
+ err = session .Wait ()
441
+ require .NoError (t , err )
442
+ })
403
443
}
404
- err := session .RequestPty ("xterm" , 128 , 128 , ssh.TerminalModes {})
405
- require .NoError (t , err )
406
- ptty := ptytest .New (t )
407
- session .Stdout = ptty .Output ()
408
- session .Stderr = ptty .Output ()
409
- session .Stdin = ptty .Input ()
410
- err = session .Start (command )
411
- require .NoError (t , err )
412
- _ = ptty .Peek (ctx , 1 ) // wait for the prompt
413
- ptty .WriteLine ("echo test" )
414
- ptty .ExpectMatch ("test" )
415
- ptty .WriteLine ("exit" )
416
- err = session .Wait ()
417
- require .NoError (t , err )
418
444
}
419
445
420
446
func TestAgent_SessionTTYExitCode (t * testing.T ) {
@@ -608,37 +634,41 @@ func TestAgent_Session_TTY_MOTD_Update(t *testing.T) {
608
634
//nolint:dogsled // Allow the blank identifiers.
609
635
conn , client , _ , _ , _ := setupAgent (t , agentsdk.Manifest {}, 0 , setSBInterval )
610
636
611
- sshClient , err := conn .SSHClient (ctx )
612
- require .NoError (t , err )
613
- t .Cleanup (func () {
614
- _ = sshClient .Close ()
615
- })
637
+ for _ , port := range sshPorts {
638
+ port := port
616
639
617
- //nolint:paralleltest // These tests need to swap the banner func.
618
- for i , test := range tests {
619
- test := test
620
- t .Run (fmt .Sprintf ("%d" , i ), func (t * testing.T ) {
621
- // Set new banner func and wait for the agent to call it to update the
622
- // banner.
623
- ready := make (chan struct {}, 2 )
624
- client .SetAnnouncementBannersFunc (func () ([]codersdk.BannerConfig , error ) {
625
- select {
626
- case ready <- struct {}{}:
627
- default :
628
- }
629
- return []codersdk.BannerConfig {test .banner }, nil
630
- })
631
- <- ready
632
- <- ready // Wait for two updates to ensure the value has propagated.
640
+ sshClient , err := conn .SSHClientOnPort (ctx , port )
641
+ require .NoError (t , err )
642
+ t .Cleanup (func () {
643
+ _ = sshClient .Close ()
644
+ })
633
645
634
- session , err := sshClient .NewSession ()
635
- require .NoError (t , err )
636
- t .Cleanup (func () {
637
- _ = session .Close ()
638
- })
646
+ //nolint:paralleltest // These tests need to swap the banner func.
647
+ for i , test := range tests {
648
+ test := test
649
+ t .Run (fmt .Sprintf ("(:%d)/%d" , port , i ), func (t * testing.T ) {
650
+ // Set new banner func and wait for the agent to call it to update the
651
+ // banner.
652
+ ready := make (chan struct {}, 2 )
653
+ client .SetAnnouncementBannersFunc (func () ([]codersdk.BannerConfig , error ) {
654
+ select {
655
+ case ready <- struct {}{}:
656
+ default :
657
+ }
658
+ return []codersdk.BannerConfig {test .banner }, nil
659
+ })
660
+ <- ready
661
+ <- ready // Wait for two updates to ensure the value has propagated.
639
662
640
- testSessionOutput (t , session , test .expected , test .unexpected , nil )
641
- })
663
+ session , err := sshClient .NewSession ()
664
+ require .NoError (t , err )
665
+ t .Cleanup (func () {
666
+ _ = session .Close ()
667
+ })
668
+
669
+ testSessionOutput (t , session , test .expected , test .unexpected , nil )
670
+ })
671
+ }
642
672
}
643
673
}
644
674
@@ -2424,6 +2454,17 @@ func setupSSHSession(
2424
2454
banner codersdk.BannerConfig ,
2425
2455
prepareFS func (fs afero.Fs ),
2426
2456
opts ... func (* agenttest.Client , * agent.Options ),
2457
+ ) * ssh.Session {
2458
+ return setupSSHSessionOnPort (t , manifest , banner , prepareFS , workspacesdk .AgentSSHPort , opts ... )
2459
+ }
2460
+
2461
+ func setupSSHSessionOnPort (
2462
+ t * testing.T ,
2463
+ manifest agentsdk.Manifest ,
2464
+ banner codersdk.BannerConfig ,
2465
+ prepareFS func (fs afero.Fs ),
2466
+ port uint16 ,
2467
+ opts ... func (* agenttest.Client , * agent.Options ),
2427
2468
) * ssh.Session {
2428
2469
ctx , cancel := context .WithTimeout (context .Background (), testutil .WaitLong )
2429
2470
defer cancel ()
@@ -2437,7 +2478,7 @@ func setupSSHSession(
2437
2478
if prepareFS != nil {
2438
2479
prepareFS (fs )
2439
2480
}
2440
- sshClient , err := conn .SSHClient (ctx )
2481
+ sshClient , err := conn .SSHClientOnPort (ctx , port )
2441
2482
require .NoError (t , err )
2442
2483
t .Cleanup (func () {
2443
2484
_ = sshClient .Close ()
0 commit comments