test(agent): improve TestAgent_Session_TTY_MOTD_Update (#10385)

This commit is contained in:
Mathias Fredriksson
2023-10-23 17:32:28 +00:00
committed by GitHub
parent 09f7b8e88c
commit 1286904de8
+30 -24
View File
@@ -380,6 +380,7 @@ func TestAgent_Session_TTY_MOTD(t *testing.T) {
} }
} }
//nolint:tparallel // Sub tests need to run sequentially.
func TestAgent_Session_TTY_MOTD_Update(t *testing.T) { func TestAgent_Session_TTY_MOTD_Update(t *testing.T) {
t.Parallel() t.Parallel()
if runtime.GOOS == "windows" { if runtime.GOOS == "windows" {
@@ -439,33 +440,38 @@ func TestAgent_Session_TTY_MOTD_Update(t *testing.T) {
} }
//nolint:dogsled // Allow the blank identifiers. //nolint:dogsled // Allow the blank identifiers.
conn, client, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, setSBInterval) conn, client, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, setSBInterval)
for _, test := range tests {
sshClient, err := conn.SSHClient(ctx)
require.NoError(t, err)
t.Cleanup(func() {
_ = sshClient.Close()
})
//nolint:paralleltest // These tests need to swap the banner func.
for i, test := range tests {
test := test test := test
// Set new banner func and wait for the agent to call it to update the t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
// banner. // Set new banner func and wait for the agent to call it to update the
ready := make(chan struct{}, 2) // banner.
client.SetServiceBannerFunc(func() (codersdk.ServiceBannerConfig, error) { ready := make(chan struct{}, 2)
select { client.SetServiceBannerFunc(func() (codersdk.ServiceBannerConfig, error) {
case ready <- struct{}{}: select {
default: case ready <- struct{}{}:
} default:
return test.banner, nil }
}) return test.banner, nil
<-ready })
<-ready // Wait for two updates to ensure the value has propagated. <-ready
<-ready // Wait for two updates to ensure the value has propagated.
sshClient, err := conn.SSHClient(ctx) session, err := sshClient.NewSession()
require.NoError(t, err) require.NoError(t, err)
t.Cleanup(func() { t.Cleanup(func() {
_ = sshClient.Close() _ = session.Close()
}) })
session, err := sshClient.NewSession()
require.NoError(t, err)
t.Cleanup(func() {
_ = session.Close()
})
testSessionOutput(t, session, test.expected, test.unexpected, nil) testSessionOutput(t, session, test.expected, test.unexpected, nil)
})
} }
} }