@@ -65,38 +65,48 @@ func TestMain(m *testing.M) {
6565 goleak .VerifyTestMain (m , testutil .GoleakOptions ... )
6666}
6767
68+ var sshPorts = []uint16 {workspacesdk .AgentSSHPort , workspacesdk .AgentStandardSSHPort }
69+
6870// NOTE: These tests only work when your default shell is bash for some reason.
6971
7072func TestAgent_Stats_SSH (t * testing.T ) {
7173 t .Parallel ()
72- ctx , cancel := context .WithTimeout (context .Background (), testutil .WaitLong )
73- defer cancel ()
7474
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 ()
7779
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 ()
8882
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+ }
100110}
101111
102112func TestAgent_Stats_ReconnectingPTY (t * testing.T ) {
@@ -278,15 +288,23 @@ func TestAgent_Stats_Magic(t *testing.T) {
278288
279289func TestAgent_SessionExec (t * testing.T ) {
280290 t .Parallel ()
281- session := setupSSHSession (t , agentsdk.Manifest {}, codersdk.ServiceBannerConfig {}, nil )
282291
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+ })
286307 }
287- output , err := session .Output (command )
288- require .NoError (t , err )
289- require .Equal (t , "test" , strings .TrimSpace (string (output )))
290308}
291309
292310//nolint:tparallel // Sub tests need to run sequentially.
@@ -396,25 +414,33 @@ func TestAgent_SessionTTYShell(t *testing.T) {
396414 // it seems like it could be either.
397415 t .Skip ("ConPTY appears to be inconsistent on Windows." )
398416 }
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+ })
403443 }
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 )
418444}
419445
420446func TestAgent_SessionTTYExitCode (t * testing.T ) {
@@ -608,37 +634,41 @@ func TestAgent_Session_TTY_MOTD_Update(t *testing.T) {
608634 //nolint:dogsled // Allow the blank identifiers.
609635 conn , client , _ , _ , _ := setupAgent (t , agentsdk.Manifest {}, 0 , setSBInterval )
610636
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
616639
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+ })
633645
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.
639662
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+ }
642672 }
643673}
644674
@@ -2424,6 +2454,17 @@ func setupSSHSession(
24242454 banner codersdk.BannerConfig ,
24252455 prepareFS func (fs afero.Fs ),
24262456 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 ),
24272468) * ssh.Session {
24282469 ctx , cancel := context .WithTimeout (context .Background (), testutil .WaitLong )
24292470 defer cancel ()
@@ -2437,7 +2478,7 @@ func setupSSHSession(
24372478 if prepareFS != nil {
24382479 prepareFS (fs )
24392480 }
2440- sshClient , err := conn .SSHClient (ctx )
2481+ sshClient , err := conn .SSHClientOnPort (ctx , port )
24412482 require .NoError (t , err )
24422483 t .Cleanup (func () {
24432484 _ = sshClient .Close ()
0 commit comments