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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy