Skip to content

Commit 51c4128

Browse files
mafredrimafredri
committed
refactor: Share setup between SSH tests, sync goroutines
1 parent 3df1a08 commit 51c4128

File tree

1 file changed

+80
-68
lines changed

1 file changed

+80
-68
lines changed

‎cli/ssh_test.go

Lines changed: 80 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cli_test
22

33
import (
4+
"context"
45
"io"
56
"net"
67
"runtime"
@@ -14,6 +15,7 @@ import (
1415

1516
"cdr.dev/slog"
1617
"cdr.dev/slog/sloggers/slogtest"
18+
1719
"github.com/coder/coder/agent"
1820
"github.com/coder/coder/cli/clitest"
1921
"github.com/coder/coder/coderd/coderdtest"
@@ -23,49 +25,52 @@ import (
2325
"github.com/coder/coder/pty/ptytest"
2426
)
2527

28+
func setupWorkspaceForSSH(t *testing.T) (*codersdk.Client, codersdk.Workspace, string) {
29+
client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerD: true})
30+
user := coderdtest.CreateFirstUser(t, client)
31+
agentToken := uuid.NewString()
32+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
33+
Parse: echo.ParseComplete,
34+
ProvisionDryRun: echo.ProvisionComplete,
35+
Provision: []*proto.Provision_Response{{
36+
Type: &proto.Provision_Response_Complete{
37+
Complete: &proto.Provision_Complete{
38+
Resources: []*proto.Resource{{
39+
Name: "dev",
40+
Type: "google_compute_instance",
41+
Agents: []*proto.Agent{{
42+
Id: uuid.NewString(),
43+
Auth: &proto.Agent_Token{
44+
Token: agentToken,
45+
},
46+
}},
47+
}},
48+
},
49+
},
50+
}},
51+
})
52+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
53+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
54+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
55+
56+
return client, workspace, agentToken
57+
}
58+
2659
func TestSSH(t *testing.T) {
27-
t.Skip("This is causing test flakes. TODO @cian fix this")
2860
t.Parallel()
2961
t.Run("ImmediateExit", func(t *testing.T) {
3062
t.Parallel()
31-
client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerD: true})
32-
user := coderdtest.CreateFirstUser(t, client)
33-
agentToken := uuid.NewString()
34-
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
35-
Parse: echo.ParseComplete,
36-
ProvisionDryRun: echo.ProvisionComplete,
37-
Provision: []*proto.Provision_Response{{
38-
Type: &proto.Provision_Response_Complete{
39-
Complete: &proto.Provision_Complete{
40-
Resources: []*proto.Resource{{
41-
Name: "dev",
42-
Type: "google_compute_instance",
43-
Agents: []*proto.Agent{{
44-
Id: uuid.NewString(),
45-
Auth: &proto.Agent_Token{
46-
Token: agentToken,
47-
},
48-
}},
49-
}},
50-
},
51-
},
52-
}},
53-
})
54-
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
55-
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
56-
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
63+
client, workspace, agentToken := setupWorkspaceForSSH(t)
5764
cmd, root := clitest.New(t, "ssh", workspace.Name)
5865
clitest.SetupConfig(t, client, root)
59-
doneChan := make(chan struct{})
6066
pty := ptytest.New(t)
6167
cmd.SetIn(pty.Input())
6268
cmd.SetErr(pty.Output())
6369
cmd.SetOut(pty.Output())
64-
go func() {
65-
defer close(doneChan)
70+
tGo(t, func() {
6671
err := cmd.Execute()
6772
assert.NoError(t, err)
68-
}()
73+
})
6974
pty.ExpectMatch("Waiting")
7075
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
7176
agentClient := codersdk.New(client.URL)
@@ -78,37 +83,12 @@ func TestSSH(t *testing.T) {
7883
})
7984
// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
8085
pty.WriteLine("exit")
81-
<-doneChan
8286
})
8387
t.Run("Stdio", func(t *testing.T) {
8488
t.Parallel()
85-
client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerD: true})
86-
user := coderdtest.CreateFirstUser(t, client)
87-
agentToken := uuid.NewString()
88-
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
89-
Parse: echo.ParseComplete,
90-
ProvisionDryRun: echo.ProvisionComplete,
91-
Provision: []*proto.Provision_Response{{
92-
Type: &proto.Provision_Response_Complete{
93-
Complete: &proto.Provision_Complete{
94-
Resources: []*proto.Resource{{
95-
Name: "dev",
96-
Type: "google_compute_instance",
97-
Agents: []*proto.Agent{{
98-
Id: uuid.NewString(),
99-
Auth: &proto.Agent_Token{
100-
Token: agentToken,
101-
},
102-
}},
103-
}},
104-
},
105-
},
106-
}},
107-
})
108-
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
109-
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
110-
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
111-
go func() {
89+
client, workspace, agentToken := setupWorkspaceForSSH(t)
90+
91+
tGoContext(t, func(ctx context.Context) {
11292
// Run this async so the SSH command has to wait for
11393
// the build and agent to connect!
11494
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
@@ -117,25 +97,22 @@ func TestSSH(t *testing.T) {
11797
agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &agent.Options{
11898
Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
11999
})
120-
t.Cleanup(func() {
121-
_ = agentCloser.Close()
122-
})
123-
}()
100+
<-ctx.Done()
101+
_ = agentCloser.Close()
102+
})
124103

125104
clientOutput, clientInput := io.Pipe()
126105
serverOutput, serverInput := io.Pipe()
127106

128107
cmd, root := clitest.New(t, "ssh", "--stdio", workspace.Name)
129108
clitest.SetupConfig(t, client, root)
130-
doneChan := make(chan struct{})
131109
cmd.SetIn(clientOutput)
132110
cmd.SetOut(serverInput)
133111
cmd.SetErr(io.Discard)
134-
go func() {
135-
defer close(doneChan)
112+
tGo(t, func() {
136113
err := cmd.Execute()
137114
assert.NoError(t, err)
138-
}()
115+
})
139116

140117
conn, channels, requests, err := ssh.NewClientConn(&stdioConn{
141118
Reader: serverOutput,
@@ -157,10 +134,45 @@ func TestSSH(t *testing.T) {
157134
err = sshClient.Close()
158135
require.NoError(t, err)
159136
_ = clientOutput.Close()
160-
<-doneChan
161137
})
162138
}
163139

140+
// tGoContext runs fn in a goroutine passing a context that will be
141+
// canceled on test completion and wait until fn has finished executing.
142+
//
143+
// NOTE(mafredri): This could be moved to a helper library.
144+
func tGoContext(t *testing.T, fn func(context.Context)) {
145+
t.Helper()
146+
147+
ctx, cancel := context.WithCancel(context.Background())
148+
done := make(chan struct{})
149+
t.Cleanup(func() {
150+
cancel()
151+
<-done
152+
})
153+
go func() {
154+
fn(ctx)
155+
close(done)
156+
}()
157+
}
158+
159+
// tGo runs fn in a goroutine and waits until fn has completed before
160+
// test completion.
161+
//
162+
// NOTE(mafredri): This could be moved to a helper library.
163+
func tGo(t *testing.T, fn func()) {
164+
t.Helper()
165+
166+
done := make(chan struct{})
167+
t.Cleanup(func() {
168+
<-done
169+
})
170+
go func() {
171+
fn()
172+
close(done)
173+
}()
174+
}
175+
164176
type stdioConn struct {
165177
io.Reader
166178
io.Writer

0 commit comments

Comments
 (0)