@@ -1690,67 +1690,85 @@ func TestSSH(t *testing.T) {
16901690 }
16911691 })
16921692
1693- t .Run ("SSHHostPrefix " , func (t * testing.T ) {
1693+ t .Run ("SSHHost " , func (t * testing.T ) {
16941694 t .Parallel ()
1695- client , workspace , agentToken := setupWorkspaceForAgent (t )
1696- _ , _ = tGoContext (t , func (ctx context.Context ) {
1697- // Run this async so the SSH command has to wait for
1698- // the build and agent to connect!
1699- _ = agenttest .New (t , client .URL , agentToken )
1700- <- ctx .Done ()
1701- })
17021695
1703- clientOutput , clientInput := io .Pipe ()
1704- serverOutput , serverInput := io .Pipe ()
1705- defer func () {
1706- for _ , c := range []io.Closer {clientOutput , clientInput , serverOutput , serverInput } {
1707- _ = c .Close ()
1708- }
1709- }()
1696+ testCases := []struct {
1697+ name , hostnameFormat string
1698+ flags []string
1699+ }{
1700+ {"Prefix" , "coder.dummy.com--%s--%s" , []string {"--ssh-host-prefix" , "coder.dummy.com--" }},
1701+ {"Suffix" , "%s--%s.coder" , []string {"--hostname-suffix" , "coder" }},
1702+ {"Both" , "%s--%s.coder" , []string {"--hostname-suffix" , "coder" , "--ssh-host-prefix" , "coder.dummy.com--" }},
1703+ }
1704+ for _ , tc := range testCases {
1705+ t .Run (tc .name , func (t * testing.T ) {
1706+ t .Parallel ()
17101707
1711- ctx , cancel := context .WithTimeout (context .Background (), testutil .WaitLong )
1712- defer cancel ()
1708+ client , workspace , agentToken := setupWorkspaceForAgent (t )
1709+ _ , _ = tGoContext (t , func (ctx context.Context ) {
1710+ // Run this async so the SSH command has to wait for
1711+ // the build and agent to connect!
1712+ _ = agenttest .New (t , client .URL , agentToken )
1713+ <- ctx .Done ()
1714+ })
17131715
1714- user , err := client .User (ctx , codersdk .Me )
1715- require .NoError (t , err )
1716+ clientOutput , clientInput := io .Pipe ()
1717+ serverOutput , serverInput := io .Pipe ()
1718+ defer func () {
1719+ for _ , c := range []io.Closer {clientOutput , clientInput , serverOutput , serverInput } {
1720+ _ = c .Close ()
1721+ }
1722+ }()
17161723
1717- inv , root := clitest .New (t , "ssh" , "--stdio" , "--ssh-host-prefix" , "coder.dummy.com--" , fmt .Sprintf ("coder.dummy.com--%s--%s" , user .Username , workspace .Name ))
1718- clitest .SetupConfig (t , client , root )
1719- inv .Stdin = clientOutput
1720- inv .Stdout = serverInput
1721- inv .Stderr = io .Discard
1724+ ctx , cancel := context .WithTimeout (context .Background (), testutil .WaitLong )
1725+ defer cancel ()
17221726
1723- cmdDone := tGo (t , func () {
1724- err := inv .WithContext (ctx ).Run ()
1725- assert .NoError (t , err )
1726- })
1727+ user , err := client .User (ctx , codersdk .Me )
1728+ require .NoError (t , err )
17271729
1728- conn , channels , requests , err := ssh .NewClientConn (& stdioConn {
1729- Reader : serverOutput ,
1730- Writer : clientInput ,
1731- }, "" , & ssh.ClientConfig {
1732- // #nosec
1733- HostKeyCallback : ssh .InsecureIgnoreHostKey (),
1734- })
1735- require .NoError (t , err )
1736- defer conn .Close ()
1730+ args := []string {"ssh" , "--stdio" }
1731+ args = append (args , tc .flags ... )
1732+ args = append (args , fmt .Sprintf (tc .hostnameFormat , user .Username , workspace .Name ))
1733+ inv , root := clitest .New (t , args ... )
1734+ clitest .SetupConfig (t , client , root )
1735+ inv .Stdin = clientOutput
1736+ inv .Stdout = serverInput
1737+ inv .Stderr = io .Discard
17371738
1738- sshClient := ssh . NewClient ( conn , channels , requests )
1739- session , err := sshClient . NewSession ()
1740- require .NoError (t , err )
1741- defer session . Close ( )
1739+ cmdDone := tGo ( t , func () {
1740+ err := inv . WithContext ( ctx ). Run ()
1741+ assert .NoError (t , err )
1742+ } )
17421743
1743- command := "sh -c exit"
1744- if runtime . GOOS == "windows" {
1745- command = "cmd.exe /c exit"
1746- }
1747- err = session . Run ( command )
1748- require . NoError ( t , err )
1749- err = sshClient . Close ( )
1750- require .NoError (t , err )
1751- _ = clientOutput .Close ()
1744+ conn , channels , requests , err := ssh . NewClientConn ( & stdioConn {
1745+ Reader : serverOutput ,
1746+ Writer : clientInput ,
1747+ }, "" , & ssh. ClientConfig {
1748+ // #nosec
1749+ HostKeyCallback : ssh . InsecureIgnoreHostKey (),
1750+ } )
1751+ require .NoError (t , err )
1752+ defer conn .Close ()
17521753
1753- <- cmdDone
1754+ sshClient := ssh .NewClient (conn , channels , requests )
1755+ session , err := sshClient .NewSession ()
1756+ require .NoError (t , err )
1757+ defer session .Close ()
1758+
1759+ command := "sh -c exit"
1760+ if runtime .GOOS == "windows" {
1761+ command = "cmd.exe /c exit"
1762+ }
1763+ err = session .Run (command )
1764+ require .NoError (t , err )
1765+ err = sshClient .Close ()
1766+ require .NoError (t , err )
1767+ _ = clientOutput .Close ()
1768+
1769+ <- cmdDone
1770+ })
1771+ }
17541772 })
17551773}
17561774
0 commit comments