Skip to content

Commit d5eb3fc

Browse files
mafredrimafredri
committed
add sub agent as part of autostart integration test
1 parent 050177b commit d5eb3fc

File tree

6 files changed

+188
-38
lines changed

6 files changed

+188
-38
lines changed

‎agent/agent_test.go

Lines changed: 160 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ import (
4848
"cdr.dev/slog/sloggers/slogtest"
4949

5050
"github.com/coder/coder/v2/agent"
51+
"github.com/coder/coder/v2/agent/agentcontainers"
5152
"github.com/coder/coder/v2/agent/agentssh"
5253
"github.com/coder/coder/v2/agent/agenttest"
5354
"github.com/coder/coder/v2/agent/proto"
@@ -60,9 +61,16 @@ import (
6061
"github.com/coder/coder/v2/tailnet"
6162
"github.com/coder/coder/v2/tailnet/tailnettest"
6263
"github.com/coder/coder/v2/testutil"
64+
"github.com/coder/quartz"
6365
)
6466

6567
func TestMain(m *testing.M) {
68+
if os.Getenv("CODER_TEST_RUN_SUB_AGENT_MAIN") == "1" {
69+
// If we're running as a subagent, we don't want to run the main tests.
70+
// Instead, we just run the subagent tests.
71+
exit := runSubAgentMain()
72+
os.Exit(exit)
73+
}
6674
goleak.VerifyTestMain(m, testutil.GoleakOptions...)
6775
}
6876

@@ -1930,6 +1938,9 @@ func TestAgent_ReconnectingPTYContainer(t *testing.T) {
19301938
if os.Getenv("CODER_TEST_USE_DOCKER") != "1" {
19311939
t.Skip("Set CODER_TEST_USE_DOCKER=1 to run this test")
19321940
}
1941+
if _, err := exec.LookPath("devcontainer"); err != nil {
1942+
t.Skip("This test requires the devcontainer CLI: npm install -g @devcontainers/cli")
1943+
}
19331944

19341945
pool, err := dockertest.NewPool("")
19351946
require.NoError(t, err, "Could not connect to docker")
@@ -1955,6 +1966,9 @@ func TestAgent_ReconnectingPTYContainer(t *testing.T) {
19551966
// nolint: dogsled
19561967
conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, func(_ *agenttest.Client, o *agent.Options) {
19571968
o.ExperimentalDevcontainersEnabled = true
1969+
o.ContainerAPIOptions = append(o.ContainerAPIOptions,
1970+
agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
1971+
)
19581972
})
19591973
ctx := testutil.Context(t, testutil.WaitLong)
19601974
ac, err := conn.ReconnectingPTY(ctx, uuid.New(), 80, 80, "/bin/sh", func(arp *workspacesdk.AgentReconnectingPTYInit) {
@@ -1986,6 +2000,60 @@ func TestAgent_ReconnectingPTYContainer(t *testing.T) {
19862000
require.ErrorIs(t, tr.ReadUntil(ctx, nil), io.EOF)
19872001
}
19882002

2003+
type subAgentRequestPayload struct {
2004+
Token string `json:"token"`
2005+
Directory string `json:"directory"`
2006+
}
2007+
2008+
// runSubAgentMain is the main function for the sub-agent that connects
2009+
// to the control plane. It reads the CODER_AGENT_URL and
2010+
// CODER_AGENT_TOKEN environment variables, sends the token, and exits
2011+
// with a status code based on the response.
2012+
func runSubAgentMain() int {
2013+
url := os.Getenv("CODER_AGENT_URL")
2014+
token := os.Getenv("CODER_AGENT_TOKEN")
2015+
if url == "" || token == "" {
2016+
_, _ = fmt.Fprintln(os.Stderr, "CODER_AGENT_URL and CODER_AGENT_TOKEN must be set")
2017+
return 10
2018+
}
2019+
2020+
dir, err := os.Getwd()
2021+
if err != nil {
2022+
_, _ = fmt.Fprintf(os.Stderr, "failed to get current working directory: %v\n", err)
2023+
return 1
2024+
}
2025+
payload := subAgentRequestPayload{
2026+
Token: token,
2027+
Directory: dir,
2028+
}
2029+
b, err := json.Marshal(payload)
2030+
if err != nil {
2031+
_, _ = fmt.Fprintf(os.Stderr, "failed to marshal payload: %v\n", err)
2032+
return 1
2033+
}
2034+
2035+
req, err := http.NewRequest("POST", url, bytes.NewReader(b))
2036+
if err != nil {
2037+
_, _ = fmt.Fprintf(os.Stderr, "failed to create request: %v\n", err)
2038+
return 1
2039+
}
2040+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
2041+
defer cancel()
2042+
req = req.WithContext(ctx)
2043+
resp, err := http.DefaultClient.Do(req)
2044+
if err != nil {
2045+
_, _ = fmt.Fprintf(os.Stderr, "agent connection failed: %v\n", err)
2046+
return 11
2047+
}
2048+
defer resp.Body.Close()
2049+
if resp.StatusCode != http.StatusOK {
2050+
_, _ = fmt.Fprintf(os.Stderr, "agent exiting with non-zero exit code %d\n", resp.StatusCode)
2051+
return 12
2052+
}
2053+
_, _ = fmt.Println("sub-agent connected successfully")
2054+
return 0
2055+
}
2056+
19892057
// This tests end-to-end functionality of auto-starting a devcontainer.
19902058
// It runs "devcontainer up" which creates a real Docker container. As
19912059
// such, it does not run by default in CI.
@@ -1999,6 +2067,56 @@ func TestAgent_DevcontainerAutostart(t *testing.T) {
19992067
if os.Getenv("CODER_TEST_USE_DOCKER") != "1" {
20002068
t.Skip("Set CODER_TEST_USE_DOCKER=1 to run this test")
20012069
}
2070+
if _, err := exec.LookPath("devcontainer"); err != nil {
2071+
t.Skip("This test requires the devcontainer CLI: npm install -g @devcontainers/cli")
2072+
}
2073+
2074+
// This HTTP handler handles requests from runSubAgentMain which
2075+
// acts as a fake sub-agent. We want to verify that the sub-agent
2076+
// connects and sends its token. We use a channel to signal
2077+
// that the sub-agent has connected successfully and then we wait
2078+
// until we receive another signal to return from the handler. This
2079+
// keeps the agent "alive" for as long as we want.
2080+
subAgentConnected := make(chan subAgentRequestPayload, 1)
2081+
subAgentReady := make(chan struct{}, 1)
2082+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
2083+
t.Logf("Sub-agent request received: %s %s", r.Method, r.URL.Path)
2084+
2085+
if r.Method != http.MethodPost {
2086+
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
2087+
return
2088+
}
2089+
2090+
// Read the token from the request body.
2091+
var payload subAgentRequestPayload
2092+
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
2093+
http.Error(w, "Failed to read token", http.StatusBadRequest)
2094+
t.Logf("Failed to read token: %v", err)
2095+
return
2096+
}
2097+
defer r.Body.Close()
2098+
2099+
t.Logf("Sub-agent request payload received: %+v", payload)
2100+
2101+
// Signal that the sub-agent has connected successfully.
2102+
select {
2103+
case <-t.Context().Done():
2104+
t.Logf("Test context done, not processing sub-agent request")
2105+
return
2106+
case subAgentConnected <- payload:
2107+
}
2108+
2109+
// Wait for the signal to return from the handler.
2110+
select {
2111+
case <-t.Context().Done():
2112+
t.Logf("Test context done, not waiting for sub-agent ready")
2113+
return
2114+
case <-subAgentReady:
2115+
}
2116+
2117+
w.WriteHeader(http.StatusOK)
2118+
}))
2119+
defer srv.Close()
20022120

20032121
pool, err := dockertest.NewPool("")
20042122
require.NoError(t, err, "Could not connect to docker")
@@ -2016,9 +2134,10 @@ func TestAgent_DevcontainerAutostart(t *testing.T) {
20162134
require.NoError(t, err, "create devcontainer directory")
20172135
devcontainerFile := filepath.Join(devcontainerPath, "devcontainer.json")
20182136
err = os.WriteFile(devcontainerFile, []byte(`{
2019-
"name": "mywork",
2020-
"image": "busybox:latest",
2021-
"cmd": ["sleep", "infinity"]
2137+
"name": "mywork",
2138+
"image": "ubuntu:latest",
2139+
"cmd": ["sleep", "infinity"],
2140+
"runArgs": ["--network=host"]
20222141
}`), 0o600)
20232142
require.NoError(t, err, "write devcontainer.json")
20242143

@@ -2043,9 +2162,24 @@ func TestAgent_DevcontainerAutostart(t *testing.T) {
20432162
},
20442163
},
20452164
}
2165+
mClock := quartz.NewMock(t)
2166+
mClock.Set(time.Now())
2167+
tickerFuncTrap := mClock.Trap().TickerFunc("agentcontainers")
2168+
20462169
//nolint:dogsled
2047-
conn, _, _, _, _ := setupAgent(t, manifest, 0, func(_ *agenttest.Client, o *agent.Options) {
2170+
_, agentClient, _, _, _ := setupAgent(t, manifest, 0, func(_ *agenttest.Client, o *agent.Options) {
20482171
o.ExperimentalDevcontainersEnabled = true
2172+
o.ContainerAPIOptions = append(
2173+
o.ContainerAPIOptions,
2174+
// Only match this specific dev container.
2175+
agentcontainers.WithClock(mClock),
2176+
agentcontainers.WithContainerLabelIncludeFilter("devcontainer.local_folder", tempWorkspaceFolder),
2177+
agentcontainers.WithSubAgentURL(srv.URL),
2178+
// The agent will copy "itself", but in the case of this test, the
2179+
// agent is actually this test binary. So we'll tell the test binary
2180+
// to execute the sub-agent main function via this env.
2181+
agentcontainers.WithSubAgentEnv("CODER_TEST_RUN_SUB_AGENT_MAIN=1"),
2182+
)
20492183
})
20502184

20512185
t.Logf("Waiting for container with label: devcontainer.local_folder=%s", tempWorkspaceFolder)
@@ -2089,32 +2223,30 @@ func TestAgent_DevcontainerAutostart(t *testing.T) {
20892223

20902224
ctx := testutil.Context(t, testutil.WaitLong)
20912225

2092-
ac, err := conn.ReconnectingPTY(ctx, uuid.New(), 80, 80, "", func(opts *workspacesdk.AgentReconnectingPTYInit) {
2093-
opts.Container = container.ID
2094-
})
2095-
require.NoError(t, err, "failed to create ReconnectingPTY")
2096-
defer ac.Close()
2097-
2098-
// Use terminal reader so we can see output in case somethin goes wrong.
2099-
tr := testutil.NewTerminalReader(t, ac)
2226+
// Ensure the container update routine runs.
2227+
tickerFuncTrap.MustWait(ctx).MustRelease(ctx)
2228+
tickerFuncTrap.Close()
2229+
_, next := mClock.AdvanceNext()
2230+
next.MustWait(ctx)
21002231

2101-
require.NoError(t, tr.ReadUntil(ctx, func(line string) bool {
2102-
return strings.Contains(line, "#") || strings.Contains(line, "$")
2103-
}), "find prompt")
2232+
// Verify that a subagent was created.
2233+
subAgents := agentClient.GetSubAgents()
2234+
require.Len(t, subAgents, 1, "expected one sub agent")
21042235

2105-
wantFileName := "file-from-devcontainer"
2106-
wantFile := filepath.Join(tempWorkspaceFolder, wantFileName)
2236+
subAgent := subAgents[0]
2237+
subAgentID, err := uuid.FromBytes(subAgent.GetId())
2238+
require.NoError(t, err, "failed to parse sub-agent ID")
2239+
t.Logf("Connecting to sub-agent: %s (ID: %s)", subAgent.Name, subAgentID)
21072240

2108-
require.NoError(t, json.NewEncoder(ac).Encode(workspacesdk.ReconnectingPTYRequest{
2109-
// NOTE(mafredri): We must use absolute path here for some reason.
2110-
Data: fmt.Sprintf("touch /workspaces/mywork/%s; exit\r", wantFileName),
2111-
}), "create file inside devcontainer")
2241+
subAgentToken, err := uuid.FromBytes(subAgent.GetAuthToken())
2242+
require.NoError(t, err, "failed to parse sub-agent token")
21122243

2113-
// Wait for the connection to close to ensure the touch was executed.
2114-
require.ErrorIs(t, tr.ReadUntil(ctx, nil), io.EOF)
2244+
payload := testutil.RequireReceive(ctx, t, subAgentConnected)
2245+
require.Equal(t, subAgentToken.String(), payload.Token, "sub-agent token should match")
2246+
require.Equal(t, "/workspaces/mywork", payload.Directory, "sub-agent directory should match")
21152247

2116-
_, err = os.Stat(wantFile)
2117-
require.NoError(t, err, "file should exist outside devcontainer")
2248+
// Allow the subagent to exit.
2249+
close(subAgentReady)
21182250
}
21192251

21202252
// TestAgent_DevcontainerRecreate tests that RecreateDevcontainer
@@ -2173,6 +2305,9 @@ func TestAgent_DevcontainerRecreate(t *testing.T) {
21732305
//nolint:dogsled
21742306
conn, client, _, _, _ := setupAgent(t, manifest, 0, func(_ *agenttest.Client, o *agent.Options) {
21752307
o.ExperimentalDevcontainersEnabled = true
2308+
o.ContainerAPIOptions = append(o.ContainerAPIOptions,
2309+
agentcontainers.WithContainerLabelIncludeFilter("devcontainer.local_folder", workspaceFolder),
2310+
)
21762311
})
21772312

21782313
ctx := testutil.Context(t, testutil.WaitLong)

‎agent/agentcontainers/api_test.go

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -302,7 +302,6 @@ func TestAPI(t *testing.T) {
302302
initialData: initialDataPayload{makeResponse(), nil},
303303
setupMock: func(mcl *acmock.MockContainerCLI, preReq *gomock.Call) {
304304
mcl.EXPECT().List(gomock.Any()).Return(makeResponse(fakeCt), nil).After(preReq).AnyTimes()
305-
mcl.EXPECT().DetectArchitecture(gomock.Any(), gomock.Any()).Return("<none>", nil).AnyTimes()
306305
},
307306
expected: makeResponse(fakeCt),
308307
},
@@ -321,7 +320,6 @@ func TestAPI(t *testing.T) {
321320
initialData: initialDataPayload{makeResponse(), assert.AnError},
322321
setupMock: func(mcl *acmock.MockContainerCLI, preReq *gomock.Call) {
323322
mcl.EXPECT().List(gomock.Any()).Return(makeResponse(fakeCt), nil).After(preReq).AnyTimes()
324-
mcl.EXPECT().DetectArchitecture(gomock.Any(), gomock.Any()).Return("<none>", nil).AnyTimes()
325323
},
326324
expected: makeResponse(fakeCt),
327325
},
@@ -338,7 +336,6 @@ func TestAPI(t *testing.T) {
338336
initialData: initialDataPayload{makeResponse(fakeCt), nil},
339337
setupMock: func(mcl *acmock.MockContainerCLI, preReq *gomock.Call) {
340338
mcl.EXPECT().List(gomock.Any()).Return(makeResponse(fakeCt2), nil).After(preReq).AnyTimes()
341-
mcl.EXPECT().DetectArchitecture(gomock.Any(), gomock.Any()).Return("<none>", nil).AnyTimes()
342339
},
343340
expected: makeResponse(fakeCt2),
344341
},
@@ -365,6 +362,7 @@ func TestAPI(t *testing.T) {
365362
api := agentcontainers.NewAPI(logger,
366363
agentcontainers.WithClock(mClock),
367364
agentcontainers.WithContainerCLI(mLister),
365+
agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
368366
)
369367
defer api.Close()
370368
r.Mount("/", api.Routes())

‎cli/exp_rpty_test.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"github.com/ory/dockertest/v3/docker"
1010

1111
"github.com/coder/coder/v2/agent"
12+
"github.com/coder/coder/v2/agent/agentcontainers"
1213
"github.com/coder/coder/v2/agent/agenttest"
1314
"github.com/coder/coder/v2/cli/clitest"
1415
"github.com/coder/coder/v2/coderd/coderdtest"
@@ -111,6 +112,9 @@ func TestExpRpty(t *testing.T) {
111112

112113
_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
113114
o.ExperimentalDevcontainersEnabled = true
115+
o.ContainerAPIOptions = append(o.ContainerAPIOptions,
116+
agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
117+
)
114118
})
115119
_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
116120

‎cli/open_test.go

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -327,8 +327,6 @@ func TestOpenVSCodeDevContainer(t *testing.T) {
327327
},
328328
}, nil,
329329
).AnyTimes()
330-
// DetectArchitecture always returns "<none>" for this test to disable agent injection.
331-
mccli.EXPECT().DetectArchitecture(gomock.Any(), gomock.Any()).Return("<none>", nil).AnyTimes()
332330

333331
client, workspace, agentToken := setupWorkspaceForAgent(t, func(agents []*proto.Agent) []*proto.Agent {
334332
agents[0].Directory = agentDir
@@ -339,7 +337,10 @@ func TestOpenVSCodeDevContainer(t *testing.T) {
339337

340338
_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
341339
o.ExperimentalDevcontainersEnabled = true
342-
o.ContainerAPIOptions = append(o.ContainerAPIOptions, agentcontainers.WithContainerCLI(mccli))
340+
o.ContainerAPIOptions = append(o.ContainerAPIOptions,
341+
agentcontainers.WithContainerCLI(mccli),
342+
agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
343+
)
343344
})
344345
_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
345346

@@ -504,8 +505,6 @@ func TestOpenVSCodeDevContainer_NoAgentDirectory(t *testing.T) {
504505
},
505506
}, nil,
506507
).AnyTimes()
507-
// DetectArchitecture always returns "<none>" for this test to disable agent injection.
508-
mccli.EXPECT().DetectArchitecture(gomock.Any(), gomock.Any()).Return("<none>", nil).AnyTimes()
509508

510509
client, workspace, agentToken := setupWorkspaceForAgent(t, func(agents []*proto.Agent) []*proto.Agent {
511510
agents[0].Name = agentName
@@ -515,7 +514,10 @@ func TestOpenVSCodeDevContainer_NoAgentDirectory(t *testing.T) {
515514

516515
_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
517516
o.ExperimentalDevcontainersEnabled = true
518-
o.ContainerAPIOptions = append(o.ContainerAPIOptions, agentcontainers.WithContainerCLI(mccli))
517+
o.ContainerAPIOptions = append(o.ContainerAPIOptions,
518+
agentcontainers.WithContainerCLI(mccli),
519+
agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
520+
)
519521
})
520522
_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
521523

‎cli/ssh_test.go

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2032,6 +2032,9 @@ func TestSSH_Container(t *testing.T) {
20322032

20332033
_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
20342034
o.ExperimentalDevcontainersEnabled = true
2035+
o.ContainerAPIOptions = append(o.ContainerAPIOptions,
2036+
agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
2037+
)
20352038
})
20362039
_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
20372040

@@ -2069,7 +2072,10 @@ func TestSSH_Container(t *testing.T) {
20692072
}, nil).AnyTimes()
20702073
_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
20712074
o.ExperimentalDevcontainersEnabled = true
2072-
o.ContainerAPIOptions = append(o.ContainerAPIOptions, agentcontainers.WithContainerCLI(mLister))
2075+
o.ContainerAPIOptions = append(o.ContainerAPIOptions,
2076+
agentcontainers.WithContainerCLI(mLister),
2077+
agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
2078+
)
20732079
})
20742080
_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
20752081

0 commit comments

Comments
 (0)