8000 add sub agent as part of autostart integration test · coder/coder@abe9116 · GitHub
[go: up one dir, main page]
Skip to content

Commit abe9116

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

File tree

5 files changed

+172
-38
lines changed

5 files changed

+172
-38
lines changed

‎agent/agent_test.go

Lines changed: 154 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")
@@ -1986,6 +1997,60 @@ func TestAgent_ReconnectingPTYContainer(t *testing.T) {
19861997
require.ErrorIs(t, tr.ReadUntil(ctx, nil), io.EOF)
19871998
}
19881999

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

20032118
pool, err := dockertest.NewPool("")
20042119
require.NoError(t, err, "Could not connect to docker")
@@ -2016,9 +2131,10 @@ func TestAgent_DevcontainerAutostart(t *testing.T) {
20162131
require.NoError(t, err, "create devcontainer directory")
20172132
devcontainerFile := filepath.Join(devcontainerPath, "devcontainer.json")
20182133
err = os.WriteFile(devcontainerFile, []byte(`{
2019-
"name": "mywork",
2020-
"image": "busybox:latest",
2021-
"cmd": ["sleep", "infinity"]
2134+
"name": "mywork",
2135+
"image": "ubuntu:latest",
2136+
"cmd": ["sleep", "infinity"],
2137+
"runArgs": ["--network=host"]
20222138
}`), 0o600)
20232139
require.NoError(t, err, "write devcontainer.json")
20242140

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

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

20902221
ctx := testutil.Context(t, testutil.WaitLong)
20912222

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)
2223+
// Ensure the container update routine runs.
2224+
tickerFuncTrap.MustWait(ctx).MustRelease(ctx)
2225+
tickerFuncTrap.Close()
2226+
_, next := mClock.AdvanceNext()
2227+
next.MustWait(ctx)
21002228

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

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

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")
2238+
subAgentToken, err := uuid.FromBytes(subAgent.GetAuthToken())
2239+
require.NoError(t, err, "failed to parse sub-agent token")
21122240

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

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

21202249
// TestAgent_DevcontainerRecreate tests that RecreateDevcontainer

‎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/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: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2069,7 +2069,10 @@ func TestSSH_Container(t *testing.T) {
20692069
}, nil).AnyTimes()
20702070
_ = agenttest.New(t, client.URL, agentToken, func(o *agent.Options) {
20712071
o.ExperimentalDevcontainersEnabled = true
2072-
o.ContainerAPIOptions = append(o.ContainerAPIOptions, agentcontainers.WithContainerCLI(mLister))
2072+
o.ContainerAPIOptions = append(o.ContainerAPIOptions,
2073+
agentcontainers.WithContainerCLI(mLister),
2074+
agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
2075+
)
20732076
})
20742077
_ = coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
20752078

‎coderd/workspaceagents_test.go

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1358,7 +1358,10 @@ func TestWorkspaceAgentContainers(t *testing.T) {
13581358
_ = agenttest.New(t, client.URL, r.AgentToken, func(o *agent.Options) {
13591359
o.Logger = logger.Named("agent")
13601360
o.ExperimentalDevcontainersEnabled = true
1361-
o.ContainerAPIOptions = append(o.ContainerAPIOptions, agentcontainers.WithContainerCLI(mcl))
1361+
o.ContainerAPIOptions = append(o.ContainerAPIOptions,
1362+
agentcontainers.WithContainerCLI(mcl),
1363+
agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
1364+
)
13621365
})
13631366
resources := coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).Wait()
13641367
require.Len(t, resources, 1, "expected one resource")
@@ -1428,8 +1431,6 @@ func TestWorkspaceAgentRecreateDevcontainer(t *testing.T) {
14281431
mccli.EXPECT().List(gomock.Any()).Return(codersdk.WorkspaceAgentListContainersResponse{
14291432
Containers: []codersdk.WorkspaceAgentContainer{devContainer},
14301433
}, nil).AnyTimes()
1431-
// DetectArchitecture always returns "<none>" for this test to disable agent injection.
1432-
mccli.EXPECT().DetectArchitecture(gomock.Any(), devContainer.ID).Return("<none>", nil).AnyTimes()
14331434
mdccli.EXPECT().Up(gomock.Any(), workspaceFolder, configFile, gomock.Any()).Return("someid", nil).Times(1)
14341435
return 0
14351436
},
@@ -1477,6 +1478,7 @@ func TestWorkspaceAgentRecreateDevcontainer(t *testing.T) {
14771478
agentcontainers.WithContainerCLI(mccli),
14781479
agentcontainers.WithDevcontainerCLI(mdccli),
14791480
agentcontainers.WithWatcher(watcher.NewNoop()),
1481+
agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
14801482
)
14811483
})
14821484
resources := coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).Wait()

0 commit comments

Comments
 (0)
0