Skip to content

Commit 6955f3f

Browse files
ethanndicksonethanndickson
committed
review p2
1 parent 726150f commit 6955f3f

File tree

5 files changed

+319
-111
lines changed

5 files changed

+319
-111
lines changed

‎coderd/database/dbfake/dbfake.go

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,20 @@ func (b WorkspaceBuildBuilder) WithAgent(mutations ...func([]*sdkproto.Agent) []
105105
Type: "aws_instance",
106106
Agents: agents,
107107
})
108+
if b.ps != nil {
109+
for _, agent := range agents {
110+
uid, err := uuid.Parse(agent.Id)
111+
require.NoError(b.t, err)
112+
msg, err := json.Marshal(wspubsub.WorkspaceEvent{
113+
Kind: wspubsub.WorkspaceEventKindAgentConnectionUpdate,
114+
WorkspaceID: b.ws.ID,
115+
AgentID: &uid,
116+
})
117+
require.NoError(b.t, err)
118+
err = b.ps.Publish(wspubsub.WorkspaceEventChannel(b.ws.OwnerID), msg)
119+
require.NoError(b.t, err)
120+
}
121+
}
108122
return b
109123
}
110124

@@ -224,6 +238,14 @@ func (b WorkspaceBuildBuilder) Do() WorkspaceResponse {
224238
}
225239
_ = dbgen.WorkspaceBuildParameters(b.t, b.db, b.params)
226240

241+
if b.ws.Deleted {
242+
err = b.db.UpdateWorkspaceDeletedByID(ownerCtx, database.UpdateWorkspaceDeletedByIDParams{
243+
ID: b.ws.ID,
244+
Deleted: true,
245+
})
246+
require.NoError(b.t, err)
247+
}
248+
227249
if b.ps != nil {
228250
msg, err := json.Marshal(wspubsub.WorkspaceEvent{
229251
Kind: wspubsub.WorkspaceEventKindStateChange,

‎coderd/workspaceagents_test.go

Lines changed: 124 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
package coderd_test
22

33
import (
4-
"bytes"
54
"context"
65
"encoding/json"
76
"fmt"
7+
"maps"
88
"net"
99
"net/http"
1010
"runtime"
@@ -1937,17 +1937,21 @@ func TestOwnedWorkspacesCoordinate(t *testing.T) {
19371937

19381938
ctx := testutil.Context(t, testutil.WaitLong)
19391939
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
1940-
firstClient, closer, _ := coderdtest.NewWithAPI(t, &coderdtest.Options{
1940+
firstClient, closer, api := coderdtest.NewWithAPI(t, &coderdtest.Options{
19411941
Coordinator: tailnet.NewCoordinator(logger),
19421942
IncludeProvisionerDaemon: true,
19431943
})
1944-
defer closer.Close()
1944+
t.Cleanup(func() {
1945+
_ = closer.Close()
1946+
})
19451947
firstUser := coderdtest.CreateFirstUser(t, firstClient)
19461948
member, memberUser := coderdtest.CreateAnotherUser(t, firstClient, firstUser.OrganizationID, rbac.RoleTemplateAdmin())
19471949

1948-
// Create a workspace
1949-
token := uuid.NewString()
1950-
resources, _ := buildWorkspaceWithAgent(t, member, firstUser.OrganizationID, token)
1950+
// Create a workspace with an agent
1951+
dbfake.WorkspaceBuild(t, api.Database, database.Workspace{
1952+
OrganizationID: firstUser.OrganizationID,
1953+
OwnerID: memberUser.ID,
1954+
}).WithAgent().Do()
19511955

19521956
u, err := member.URL.Parse("/api/v2/tailnet")
19531957
require.NoError(t, err)
@@ -1990,69 +1994,40 @@ func TestOwnedWorkspacesCoordinate(t *testing.T) {
19901994
// Existing agent
19911995
require.Len(t, update.UpsertedAgents, 1)
19921996
require.Equal(t, update.UpsertedAgents[0].WorkspaceId, wsID)
1993-
require.EqualValues(t, update.UpsertedAgents[0].Id, resources[0].Agents[0].ID)
19941997

19951998
require.Len(t, update.DeletedWorkspaces, 0)
19961999
require.Len(t, update.DeletedAgents, 0)
19972000

19982001
// Build a second workspace
1999-
secondToken := uuid.NewString()
2000-
secondResources, secondWorkspace := buildWorkspaceWithAgent(t, member, firstUser.OrganizationID, secondToken)
2001-
2002-
// Workspace starting
2003-
update, err = stream.Recv()
2004-
require.NoError(t, err)
2005-
require.Len(t, update.UpsertedWorkspaces, 1)
2006-
require.Equal(t, update.UpsertedWorkspaces[0].Status, tailnetproto.Workspace_STARTING)
2007-
2008-
require.Len(t, update.DeletedWorkspaces, 0)
2009-
require.Len(t, update.DeletedAgents, 0)
2010-
require.Len(t, update.UpsertedAgents, 0)
2011-
2012-
// Workspace running, agent created
2013-
update, err = stream.Recv()
2014-
require.NoError(t, err)
2015-
require.Len(t, update.UpsertedWorkspaces, 1)
2016-
require.Equal(t, update.UpsertedWorkspaces[0].Status, tailnetproto.Workspace_RUNNING)
2017-
wsID = update.UpsertedWorkspaces[0].Id
2018-
require.Len(t, update.UpsertedAgents, 1)
2019-
require.Equal(t, update.UpsertedAgents[0].WorkspaceId, wsID)
2020-
require.EqualValues(t, update.UpsertedAgents[0].Id, secondResources[0].Agents[0].ID)
2021-
2022-
require.Len(t, update.DeletedWorkspaces, 0)
2023-
require.Len(t, update.DeletedAgents, 0)
2024-
2025-
_, err = member.CreateWorkspaceBuild(ctx, secondWorkspace.ID, codersdk.CreateWorkspaceBuildRequest{
2026-
Transition: codersdk.WorkspaceTransitionDelete,
2027-
})
2028-
require.NoError(t, err)
2029-
2030-
// Wait for the workspace to be deleted
2031-
deletedAgents := make([]*tailnetproto.Agent, 0)
2032-
workspaceUpdates := make([]*tailnetproto.Workspace, 0)
2033-
require.Eventually(t, func() bool {
2034-
update, err = stream.Recv()
2035-
if err != nil {
2036-
return false
2037-
}
2038-
deletedAgents = append(deletedAgents, update.DeletedAgents...)
2039-
workspaceUpdates = append(workspaceUpdates, update.UpsertedWorkspaces...)
2040-
return len(update.DeletedWorkspaces) == 1 &&
2041-
bytes.Equal(update.DeletedWorkspaces[0].Id, wsID)
2042-
}, testutil.WaitMedium, testutil.IntervalSlow)
2043-
2044-
// We should have seen an update for the agent being deleted
2045-
require.Len(t, deletedAgents, 1)
2046-
require.EqualValues(t, deletedAgents[0].Id, secondResources[0].Agents[0].ID)
2047-
2048-
// But we may also see a 'pending' state transition before 'deleting'
2049-
deletingFound := false
2050-
for _, ws := range workspaceUpdates {
2051-
if bytes.Equal(ws.Id, wsID) && ws.Status == tailnetproto.Workspace_DELETING {
2052-
deletingFound = true
2053-
}
2002+
secondWorkspace := dbfake.WorkspaceBuild(t, api.Database, database.Workspace{
2003+
OrganizationID: firstUser.OrganizationID,
2004+
OwnerID: memberUser.ID,
2005+
}).WithAgent().Pubsub(api.Pubsub).Do()
2006+
2007+
// Wait for the second workspace to be running with an agent
2008+
expectedState := map[uuid.UUID]workspace{
2009+
secondWorkspace.Workspace.ID: {
2010+
Status: tailnetproto.Workspace_RUNNING,
2011+
NumAgents: 1,
2012+
},
20542013
}
2055-
require.True(t, deletingFound)
2014+
waitForUpdates(t, ctx, stream, map[uuid.UUID]workspace{}, expectedState)
2015+
2016+
// Wait for the workspace and agent to be deleted
2017+
secondWorkspace.Workspace.Deleted = true
2018+
dbfake.WorkspaceBuild(t, api.Database, secondWorkspace.Workspace).
2019+
Seed(database.WorkspaceBuild{
2020+
Transition: database.WorkspaceTransitionDelete,
2021+
BuildNumber: 2,
2022+
}).Pubsub(api.Pubsub).Do()
2023+
2024+
priorState := expectedState
2025+
waitForUpdates(t, ctx, stream, priorState, map[uuid.UUID]workspace{
2026+
secondWorkspace.Workspace.ID: {
2027+
Status: tailnetproto.Workspace_DELETED,
2028+
NumAgents: 0,
2029+
},
2030+
})
20562031
}
20572032

20582033
func requireGetManifest(ctx context.Context, t testing.TB, aAPI agentproto.DRPCAgentClient) agentsdk.Manifest {
@@ -2075,17 +2050,90 @@ func postStartup(ctx context.Context, t testing.TB, client agent.Client, startup
20752050
return err
20762051
}
20772052

2078-
func buildWorkspaceWithAgent(t *testing.T, client *codersdk.Client, orgID uuid.UUID, agentToken string) ([]codersdk.WorkspaceResource, codersdk.Workspace) {
2079-
version := coderdtest.CreateTemplateVersion(t, client, orgID, &echo.Responses{
2080-
Parse: echo.ParseComplete,
2081-
ProvisionPlan: echo.PlanComplete,
2082-
ProvisionApply: echo.ProvisionApplyWithAgent(agentToken),
2083-
})
2084-
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
2085-
template := coderdtest.CreateTemplate(t, client, orgID, version.ID)
2086-
workspace := coderdtest.CreateWorkspace(t, client, template.ID)
2087-
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
2088-
_ = agenttest.New(t, client.URL, agentToken)
2089-
resources := coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
2090-
return resources, workspace
2053+
type workspace struct {
2054+
Status tailnetproto.Workspace_Status
2055+
NumAgents int
2056+
}
2057+
2058+
func waitForUpdates(
2059+
t *testing.T,
2060+
//nolint:revive // t takes precedence
2061+
ctx context.Context,
2062+
stream tailnetproto.DRPCTailnet_WorkspaceUpdatesClient,
2063+
currentState map[uuid.UUID]workspace,
2064+
expectedState map[uuid.UUID]workspace,
2065+
) {
2066+
t.Helper()
2067+
errCh := make(chan error, 1)
2068+
go func() {
2069+
for {
2070+
select {
2071+
case <-ctx.Done():
2072+
errCh <- ctx.Err()
2073+
return
2074+
default:
2075+
}
2076+
update, err := stream.Recv()
2077+
if err != nil {
2078+
errCh <- err
2079+
return
2080+
}
2081+
for _, ws := range update.UpsertedWorkspaces {
2082+
id, err := uuid.FromBytes(ws.Id)
2083+
if err != nil {
2084+
errCh <- err
2085+
return
2086+
}
2087+
currentState[id] = workspace{
2088+
Status: ws.Status,
2089+
NumAgents: currentState[id].NumAgents,
2090+
}
2091+
}
2092+
for _, ws := range update.DeletedWorkspaces {
2093+
id, err := uuid.FromBytes(ws.Id)
2094+
if err != nil {
2095+
errCh <- err
2096+
return
2097+
}
2098+
currentState[id] = workspace{
2099+
Status: tailnetproto.Workspace_DELETED,
2100+
NumAgents: currentState[id].NumAgents,
2101+
}
2102+
}
2103+
for _, a := range update.UpsertedAgents {
2104+
id, err := uuid.FromBytes(a.WorkspaceId)
2105+
if err != nil {
2106+
errCh <- err
2107+
return
2108+
}
2109+
currentState[id] = workspace{
2110+
Status: currentState[id].Status,
2111+
NumAgents: currentState[id].NumAgents + 1,
2112+
}
2113+
}
2114+
for _, a := range update.DeletedAgents {
2115+
id, err := uuid.FromBytes(a.WorkspaceId)
2116+
if err != nil {
2117+
errCh <- err
2118+
return
2119+
}
2120+
currentState[id] = workspace{
2121+
Status: currentState[id].Status,
2122+
NumAgents: currentState[id].NumAgents - 1,
2123+
}
2124+
}
2125+
if maps.Equal(currentState, expectedState) {
2126+
errCh <- nil
2127+
return
2128+
}
2129+
}
2130+
}()
2131+
select {
2132+
case err := <-errCh:
2133+
if err != nil {
2134+
t.Fatal(err)
2135+
}
2136+
case <-ctx.Done():
2137+
t.Fatal("Timeout waiting for desired state")
2138+
}
20912139
}

0 commit comments

Comments
 (0)