Skip to content

Commit df424e6

Browse files
spikecurtisspikecurtis
committed
agent_test running SSH in pty use ptytest.Start
Signed-off-by: Spike Curtis <spike@coder.com>
1 parent 8ec3d1f commit df424e6

File tree

1 file changed

+17
-42
lines changed

1 file changed

+17
-42
lines changed

‎agent/agent_test.go

Lines changed: 17 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ import (
2424
"testing"
2525
"time"
2626

27+
"github.com/coder/coder/pty"
28+
2729
scp "github.com/bramvdbogaerde/go-scp"
2830
"github.com/google/uuid"
2931
"github.com/pion/udp"
@@ -481,17 +483,10 @@ func TestAgent_TCPLocalForwarding(t *testing.T) {
481483
}
482484
}()
483485

484-
pty := ptytest.New(t)
485-
486-
cmd := setupSSHCommand(t, []string{"-L", fmt.Sprintf("%d:127.0.0.1:%d", randomPort, remotePort)}, []string{"sleep", "5"})
487-
cmd.Stdin = pty.Input()
488-
cmd.Stdout = pty.Output()
489-
cmd.Stderr = pty.Output()
490-
err = cmd.Start()
491-
require.NoError(t, err)
486+
_, proc := setupSSHCommand(t, []string{"-L", fmt.Sprintf("%d:127.0.0.1:%d", randomPort, remotePort)}, []string{"sleep", "5"})
492487

493488
go func() {
494-
err := cmd.Wait()
489+
err := proc.Wait()
495490
select {
496491
case <-done:
497492
default:
@@ -523,7 +518,7 @@ func TestAgent_TCPLocalForwarding(t *testing.T) {
523518

524519
<-done
525520

526-
_ = cmd.Process.Kill()
521+
_ = proc.Kill()
527522
}
528523

529524
//nolint:paralleltest // This test reserves a port.
@@ -562,17 +557,10 @@ func TestAgent_TCPRemoteForwarding(t *testing.T) {
562557
}
563558
}()
564559

565-
pty := ptytest.New(t)
566-
567-
cmd := setupSSHCommand(t, []string{"-R", fmt.Sprintf("127.0.0.1:%d:127.0.0.1:%d", randomPort, localPort)}, []string{"sleep", "5"})
568-
cmd.Stdin = pty.Input()
569-
cmd.Stdout = pty.Output()
570-
cmd.Stderr = pty.Output()
571-
err = cmd.Start()
572-
require.NoError(t, err)
560+
_, proc := setupSSHCommand(t, []string{"-R", fmt.Sprintf("127.0.0.1:%d:127.0.0.1:%d", randomPort, localPort)}, []string{"sleep", "5"})
573561

574562
go func() {
575-
err := cmd.Wait()
563+
err := proc.Wait()
576564
select {
577565
case <-done:
578566
default:
@@ -604,7 +592,7 @@ func TestAgent_TCPRemoteForwarding(t *testing.T) {
604592

605593
<-done
606594

607-
_ = cmd.Process.Kill()
595+
_ = proc.Kill()
608596
}
609597

610598
func TestAgent_UnixLocalForwarding(t *testing.T) {
@@ -641,17 +629,10 @@ func TestAgent_UnixLocalForwarding(t *testing.T) {
641629
}
642630
}()
643631

644-
pty := ptytest.New(t)
645-
646-
cmd := setupSSHCommand(t, []string{"-L", fmt.Sprintf("%s:%s", localSocketPath, remoteSocketPath)}, []string{"sleep", "5"})
647-
cmd.Stdin = pty.Input()
648-
cmd.Stdout = pty.Output()
649-
cmd.Stderr = pty.Output()
650-
err = cmd.Start()
651-
require.NoError(t, err)
632+
_, proc := setupSSHCommand(t, []string{"-L", fmt.Sprintf("%s:%s", localSocketPath, remoteSocketPath)}, []string{"sleep", "5"})
652633

653634
go func() {
654-
err := cmd.Wait()
635+
err := proc.Wait()
655636
select {
656637
case <-done:
657638
default:
@@ -676,7 +657,7 @@ func TestAgent_UnixLocalForwarding(t *testing.T) {
676657
_ = conn.Close()
677658
<-done
678659

679-
_ = cmd.Process.Kill()
660+
_ = proc.Kill()
680661
}
681662

682663
func TestAgent_UnixRemoteForwarding(t *testing.T) {
@@ -713,17 +694,10 @@ func TestAgent_UnixRemoteForwarding(t *testing.T) {
713694
}
714695
}()
715696

716-
pty := ptytest.New(t)
717-
718-
cmd := setupSSHCommand(t, []string{"-R", fmt.Sprintf("%s:%s", remoteSocketPath, localSocketPath)}, []string{"sleep", "5"})
719-
cmd.Stdin = pty.Input()
720-
cmd.Stdout = pty.Output()
721-
cmd.Stderr = pty.Output()
722-
err = cmd.Start()
723-
require.NoError(t, err)
697+
_, proc := setupSSHCommand(t, []string{"-R", fmt.Sprintf("%s:%s", remoteSocketPath, localSocketPath)}, []string{"sleep", "5"})
724698

725699
go func() {
726-
err := cmd.Wait()
700+
err := proc.Wait()
727701
select {
728702
case <-done:
729703
default:
@@ -750,7 +724,7 @@ func TestAgent_UnixRemoteForwarding(t *testing.T) {
750724

751725
<-done
752726

753-
_ = cmd.Process.Kill()
727+
_ = proc.Kill()
754728
}
755729

756730
func TestAgent_SFTP(t *testing.T) {
@@ -1629,7 +1603,7 @@ func TestAgent_WriteVSCodeConfigs(t *testing.T) {
16291603
}, testutil.WaitShort, testutil.IntervalFast)
16301604
}
16311605

1632-
func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
1606+
func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) (*ptytest.PTYCmd, pty.Process) {
16331607
//nolint:dogsled
16341608
agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
16351609
listener, err := net.Listen("tcp", "127.0.0.1:0")
@@ -1671,7 +1645,8 @@ func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exe
16711645
"host",
16721646
)
16731647
args = append(args, afterArgs...)
1674-
return exec.Command("ssh", args...)
1648+
cmd := exec.Command("ssh", args...)
1649+
return ptytest.Start(t, cmd)
16751650
}
16761651

16771652
func setupSSHSession(t *testing.T, options agentsdk.Manifest) *ssh.Session {

0 commit comments

Comments
 (0)