mirror of
https://github.com/coder/coder.git
synced 2026-06-05 14:08:20 +00:00
Cian Johnston
4d0d187806
The `create_workspace` tool waited for the workspace build to succeed and the agent to become connectable, but did not wait for the agent's startup scripts (e.g. git clone) to finish. This caused agents to attempt file operations on repositories that hadn't been cloned yet. Add a waitForStartupScripts step that polls the agent's lifecycle_state via GetWorkspaceAgentLifecycleStateByID until it transitions out of created/starting into a terminal state (ready, start_error, or start_timeout). The tool now only returns success once the workspace is fully initialized. If the scripts fail or time out, the tool still returns (non-fatal) with an appropriate agent_status so the model knows something went wrong. Created using thingies (Opus 4.6 Max)
111 lines
3.2 KiB
Go
111 lines
3.2 KiB
Go
package chattool //nolint:testpackage // Uses internal symbols.
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/require"
|
|
"go.uber.org/mock/gomock"
|
|
|
|
"github.com/coder/coder/v2/coderd/database"
|
|
"github.com/coder/coder/v2/coderd/database/dbmock"
|
|
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
|
)
|
|
|
|
func TestWaitForAgentReady(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
t.Run("AgentConnectsAndLifecycleReady", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
agentID := uuid.New()
|
|
|
|
// Mock returns Ready lifecycle state.
|
|
db.EXPECT().
|
|
GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID).
|
|
Return(database.GetWorkspaceAgentLifecycleStateByIDRow{
|
|
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
|
|
}, nil)
|
|
|
|
// AgentConnFn succeeds immediately.
|
|
connFn := func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
|
return nil, func() {}, nil
|
|
}
|
|
|
|
result := waitForAgentReady(context.Background(), db, agentID, connFn)
|
|
require.Empty(t, result)
|
|
})
|
|
|
|
t.Run("AgentConnectTimeout", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
agentID := uuid.New()
|
|
|
|
// AgentConnFn always fails - context will timeout.
|
|
connFn := func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
|
return nil, nil, context.DeadlineExceeded
|
|
}
|
|
|
|
// Use a context that's already canceled to avoid waiting.
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
result := waitForAgentReady(ctx, db, agentID, connFn)
|
|
require.Equal(t, "not_ready", result["agent_status"])
|
|
require.NotEmpty(t, result["agent_error"])
|
|
})
|
|
|
|
t.Run("AgentConnectsButStartupFails", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
agentID := uuid.New()
|
|
|
|
// Mock returns StartError lifecycle state.
|
|
db.EXPECT().
|
|
GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID).
|
|
Return(database.GetWorkspaceAgentLifecycleStateByIDRow{
|
|
LifecycleState: database.WorkspaceAgentLifecycleStateStartError,
|
|
}, nil)
|
|
|
|
connFn := func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
|
return nil, func() {}, nil
|
|
}
|
|
|
|
result := waitForAgentReady(context.Background(), db, agentID, connFn)
|
|
require.Equal(t, "startup_scripts_failed", result["startup_scripts"])
|
|
require.Equal(t, "start_error", result["lifecycle_state"])
|
|
})
|
|
|
|
t.Run("NilAgentConnFn", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
agentID := uuid.New()
|
|
|
|
// Mock returns Ready lifecycle state.
|
|
db.EXPECT().
|
|
GetWorkspaceAgentLifecycleStateByID(gomock.Any(), agentID).
|
|
Return(database.GetWorkspaceAgentLifecycleStateByIDRow{
|
|
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
|
|
}, nil)
|
|
|
|
result := waitForAgentReady(context.Background(), db, agentID, nil)
|
|
require.Empty(t, result)
|
|
})
|
|
|
|
t.Run("NilDB", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
connFn := func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
|
return nil, func() {}, nil
|
|
}
|
|
|
|
result := waitForAgentReady(context.Background(), nil, uuid.New(), connFn)
|
|
require.Empty(t, result)
|
|
})
|
|
}
|