Skip to content

Commit 8f08e00

Browse files
dannykoppingdannykopping
committed
Add tests
Signed-off-by: Danny Kopping <danny@coder.com>
1 parent 5438b65 commit 8f08e00

File tree

1 file changed

+116
-21
lines changed

1 file changed

+116
-21
lines changed

‎provisionersdk/agent_test.go

Lines changed: 116 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -7,57 +7,152 @@
77
package provisionersdk_test
88

99
import (
10+
"bytes"
11+
"context"
12+
"errors"
1013
"fmt"
1114
"net/http"
1215
"net/http/httptest"
1316
"net/url"
1417
"os/exec"
1518
"runtime"
1619
"strings"
20+
"sync"
1721
"testing"
22+
"time"
1823

1924
"github.com/go-chi/render"
2025
"github.com/stretchr/testify/require"
2126

27+
"github.com/coder/coder/v2/testutil"
28+
2229
"github.com/coder/coder/v2/provisionersdk"
2330
)
2431

32+
// mimicking the --version output which we use to test the binary (see provisionersdk/scripts/bootstrap_*).
33+
const versionOutput = `Coder v2.11.0+8979bfe Tue May 7 17:30:19 UTC 2024`
34+
2535
// bashEcho is a script that calls the local `echo` with the arguments. This is preferable to
2636
// sending the real `echo` binary since macOS 14.4+ immediately sigkills `echo` if it is copied to
2737
// another directory and run locally.
2838
const bashEcho = `#!/usr/bin/env bash
29-
echo $@`
39+
echo "` + versionOutput + `"`
40+
41+
const unexpectedEcho = `#!/usr/bin/env bash
42+
echo "this is not the agent you are looking for"`
3043

3144
func TestAgentScript(t *testing.T) {
3245
t.Parallel()
33-
t.Run("Run", func(t *testing.T) {
46+
47+
t.Run("Valid", func(t *testing.T) {
3448
t.Parallel()
35-
srv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
36-
render.Status(r, http.StatusOK)
37-
render.Data(rw, r, []byte(bashEcho))
38-
}))
39-
defer srv.Close()
40-
srvURL, err := url.Parse(srv.URL)
41-
require.NoError(t, err)
4249

43-
script, exists := provisionersdk.AgentScriptEnv()[fmt.Sprintf("CODER_AGENT_SCRIPT_%s_%s", runtime.GOOS, runtime.GOARCH)]
44-
if !exists {
45-
t.Skip("Agent not supported...")
46-
return
47-
}
48-
script = strings.ReplaceAll(script, "${ACCESS_URL}", srvURL.String()+"/")
49-
script = strings.ReplaceAll(script, "${AUTH_TYPE}", "token")
50+
script := serveScript(t, bashEcho)
51+
52+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
53+
t.Cleanup(cancel)
54+
55+
var output bytes.Buffer
5056
// This is intentionally ran in single quotes to mimic how a customer may
5157
// embed our script. Our scripts should not include any single quotes.
5258
// nolint:gosec
53-
output, err := exec.Command("sh", "-c", "sh -c '"+script+"'").CombinedOutput()
54-
t.Log(string(output))
59+
cmd := exec.CommandContext(ctx, "sh", "-c", "sh -c '"+script+"'")
60+
cmd.Stdout = &output
61+
cmd.Stderr = &output
62+
require.NoError(t, cmd.Start())
63+
64+
err := cmd.Wait()
65+
if err != nil {
66+
var exitErr *exec.ExitError
67+
if errors.As(err, &exitErr) {
68+
require.Equal(t, 0, exitErr.ExitCode())
69+
} else {
70+
t.Fatalf("unexpected err: %s", err)
71+
}
72+
}
73+
74+
t.Log(output.String())
5575
require.NoError(t, err)
5676
// Ignore debug output from `set -x`, we're only interested in the last line.
57-
lines := strings.Split(strings.TrimSpace(string(output)), "\n")
77+
lines := strings.Split(strings.TrimSpace(output.String()), "\n")
5878
lastLine := lines[len(lines)-1]
59-
// Because we use the "echo" binary, we should expect the arguments provided
79+
// When we use the "bashEcho" binary, we should expect the arguments provided
6080
// as the response to executing our script.
61-
require.Equal(t, "agent", lastLine)
81+
require.Equal(t, versionOutput, lastLine)
6282
})
83+
84+
t.Run("Invalid", func(t *testing.T) {
85+
t.Parallel()
86+
87+
script := serveScript(t, unexpectedEcho)
88+
89+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
90+
t.Cleanup(cancel)
91+
92+
var output bytes.Buffer
93+
// This is intentionally ran in single quotes to mimic how a customer may
94+
// embed our script. Our scripts should not include any single quotes.
95+
// nolint:gosec
96+
cmd := exec.CommandContext(ctx, "sh", "-c", "sh -c '"+script+"'")
97+
cmd.WaitDelay = time.Second
98+
cmd.Stdout = &output
99+
cmd.Stderr = &output
100+
require.NoError(t, cmd.Start())
101+
102+
done := make(chan error, 1)
103+
var wg sync.WaitGroup
104+
wg.Add(1)
105+
go func() {
106+
defer wg.Done()
107+
108+
// The bootstrap scripts trap exit codes to allow operators to view the script logs and debug the process
109+
// while it is still running. We do not expect Wait() to complete.
110+
err := cmd.Wait()
111+
done <- err
112+
}()
113+
114+
select {
115+
case <-ctx.Done():
116+
// Timeout.
117+
break
118+
case err := <-done:
119+
// If done signals before context times out, script behaved in an unexpected way.
120+
if err != nil {
121+
t.Fatalf("unexpected err: %s", err)
122+
}
123+
}
124+
125+
// Kill the command, wait for the command to yield.
126+
require.NoError(t, cmd.Cancel())
127+
wg.Wait()
128+
129+
t.Log(output.String())
130+
131+
require.Eventually(t, func() bool {
132+
return bytes.Contains(output.Bytes(), []byte("ERROR: Downloaded agent binary is invalid"))
133+
}, testutil.WaitShort, testutil.IntervalSlow)
134+
})
135+
}
136+
137+
// serveScript creates a fake HTTP server which serves a requested "agent binary" (which is actually the given input string)
138+
// which will be attempted to run to verify that it is correct.
139+
func serveScript(t *testing.T, in string) string {
140+
t.Helper()
141+
142+
srv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
143+
render.Status(r, http.StatusOK)
144+
render.Data(rw, r, []byte(in))
145+
}))
146+
t.Cleanup(srv.Close)
147+
srvURL, err := url.Parse(srv.URL)
148+
require.NoError(t, err)
149+
150+
script, exists := provisionersdk.AgentScriptEnv()[fmt.Sprintf("CODER_AGENT_SCRIPT_%s_%s", runtime.GOOS, runtime.GOARCH)]
151+
if !exists {
152+
t.Skip("Agent not supported...")
153+
return ""
154+
}
155+
script = strings.ReplaceAll(script, "${ACCESS_URL}", srvURL.String()+"/")
156+
script = strings.ReplaceAll(script, "${AUTH_TYPE}", "token")
157+
return script
63158
}

0 commit comments

Comments
 (0)