Skip to content

Commit 056a454

Browse files
sreyasreya
committed
add tests for workspacelock
1 parent 6113f97 commit 056a454

File tree

7 files changed

+101
-121
lines changed

7 files changed

+101
-121
lines changed

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2545,8 +2545,6 @@ func (q *querier) UpdateWorkspaceTTLToBeWithinTemplateMax(ctx context.Context, a
25452545
}
25462546

25472547
func (q *querier) UpdateWorkspacesDeletingAtByTemplateID(ctx context.Context, arg database.UpdateWorkspacesDeletingAtByTemplateIDParams) error {
2548-
// TODO: This is kinda messed up. We want to update this whenever we update the locked_ttl on a template. But a template
2549-
// admin may not have permissions to edit workspaces so this won't work.
25502548
fetch := func(ctx context.Context, arg database.UpdateWorkspacesDeletingAtByTemplateIDParams) (database.Template, error) {
25512549
return q.db.GetTemplateByID(ctx, arg.TemplateID)
25522550
}

‎coderd/database/queries.sql.go

Lines changed: 9 additions & 9 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/workspaces.sql

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -491,21 +491,21 @@ WHERE
491491
-- name: UpdateWorkspaceLockedDeletingAt :exec
492492
UPDATE workspaces
493493
SET
494-
locked_at = @locked_at,
495-
-- If locked_at is null (meaning unlocked) or the template-defined locked_ttl is 0 we should set
496-
-- deleting_at to NULL else set it to the locked_at + locked_ttl duration.
497-
deleting_at = CASE WHEN @locked_at IS NULL OR @locked_ttl_ms::bigint = 0 THEN NULL ELSE '@locked_at' + interval '1 millisecond' * @locked_ttl_ms::bigint END,
494+
locked_at = $2,
498495
-- When a workspace is unlocked we want to update the last_used_at to avoid the workspace getting re-locked.
499496
-- if we're locking the workspace then we leave it alone.
500-
last_used_at = CASE WHEN @locked_at IS NULL THEN now() at time zone 'utc' ELSE last_used_at END
497+
last_used_at = CASE WHEN $2::timestamptz IS NULL THEN now() at time zone 'utc' ELSE last_used_at END,
498+
-- If locked_at is null (meaning unlocked) or the template-defined locked_ttl is 0 we should set
499+
-- deleting_at to NULL else set it to the locked_at + locked_ttl duration.
500+
deleting_at = CASE WHEN $2::timestamptz IS NULL OR @locked_ttl_ms::bigint = 0 THEN NULL ELSE $2::timestamptz + interval '1 millisecond' * @locked_ttl_ms::bigint END
501501
WHERE
502-
id = @id;
502+
id = $1;
503503

504504
-- name: UpdateWorkspacesDeletingAtByTemplateID :exec
505505
UPDATE
506506
workspaces
507507
SET
508-
deleting_at = locked_at + interval '1 milliseconds' * @locked_ttl_ms::bigint
508+
deleting_at = CASE WHEN @locked_ttl_ms::bigint = 0 THEN NULL ELSE locked_at + interval '1 milliseconds' * @locked_ttl_ms::bigint END
509509
WHERE
510510
template_id = @template_id
511511
AND

‎coderd/workspaces.go

Lines changed: 7 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ import (
1212

1313
"github.com/go-chi/chi/v5"
1414
"github.com/google/uuid"
15-
"golang.org/x/exp/slices"
1615
"golang.org/x/xerrors"
1716

1817
"cdr.dev/slog"
@@ -1120,6 +1119,11 @@ func convertWorkspace(
11201119
lockedAt = &workspace.LockedAt.Time
11211120
}
11221121

1122+
var deletedAt *time.Time
1123+
if workspace.DeletingAt.Valid {
1124+
deletedAt = &workspace.DeletingAt.Time
1125+
}
1126+
11231127
failingAgents := []uuid.UUID{}
11241128
for _, resource := range workspaceBuild.Resources {
11251129
for _, agent := range resource.Agents {
@@ -1130,8 +1134,7 @@ func convertWorkspace(
11301134
}
11311135

11321136
var (
1133-
ttlMillis = convertWorkspaceTTLMillis(workspace.Ttl)
1134-
deletingAt = calculateDeletingAt(workspace, template, workspaceBuild)
1137+
ttlMillis = convertWorkspaceTTLMillis(workspace.Ttl)
11351138
)
11361139

11371140
return codersdk.Workspace{
@@ -1152,7 +1155,7 @@ func convertWorkspace(
11521155
AutostartSchedule: autostartSchedule,
11531156
TTLMillis: ttlMillis,
11541157
LastUsedAt: workspace.LastUsedAt,
1155-
DeletingAt: deletingAt,
1158+
DeletingAt: deletedAt,
11561159
LockedAt: lockedAt,
11571160
Health: codersdk.WorkspaceHealth{
11581161
Healthy: len(failingAgents) == 0,
@@ -1170,19 +1173,6 @@ func convertWorkspaceTTLMillis(i sql.NullInt64) *int64 {
11701173
return &millis
11711174
}
11721175

1173-
// Calculate the time of the upcoming workspace deletion, if applicable; otherwise, return nil.
1174-
// Workspaces may have impending deletions if InactivityTTL feature is turned on and the workspace is inactive.
1175-
func calculateDeletingAt(workspace database.Workspace, template database.Template, build codersdk.WorkspaceBuild) *time.Time {
1176-
inactiveStatuses := []codersdk.WorkspaceStatus{codersdk.WorkspaceStatusStopped, codersdk.WorkspaceStatusCanceled, codersdk.WorkspaceStatusFailed, codersdk.WorkspaceStatusDeleted}
1177-
isInactive := slices.Contains(inactiveStatuses, build.Status)
1178-
// If InactivityTTL is turned off (set to 0) or if the workspace is active, there is no impending deletion
1179-
if template.InactivityTTL == 0 || !isInactive {
1180-
return nil
1181-
}
1182-
1183-
return ptr.Ref(workspace.LastUsedAt.Add(time.Duration(template.InactivityTTL) * time.Nanosecond))
1184-
}
1185-
11861176
func validWorkspaceTTLMillis(millis *int64, templateDefault, templateMax time.Duration) (sql.NullInt64, error) {
11871177
if templateDefault == 0 && templateMax != 0 || (templateMax > 0 && templateDefault > templateMax) {
11881178
templateDefault = templateMax

‎coderd/workspaces_internal_test.go

Lines changed: 0 additions & 82 deletions
This file was deleted.

‎coderd/workspaces_test.go

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2653,25 +2653,34 @@ func TestWorkspaceLock(t *testing.T) {
26532653
user = coderdtest.CreateFirstUser(t, client)
26542654
version = coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
26552655
_ = coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
2656-
template = coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
2657-
workspace = coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
2658-
_ = coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
2656+
lockedTTL = time.Minute
26592657
)
26602658

2659+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID, func(ctr *codersdk.CreateTemplateRequest) {
2660+
ctr.LockedTTLMillis = ptr.Ref[int64](lockedTTL.Milliseconds())
2661+
})
2662+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
2663+
_ = coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
2664+
26612665
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
26622666
defer cancel()
26632667

2668+
lastUsedAt := workspace.LastUsedAt
26642669
err := client.UpdateWorkspaceLock(ctx, workspace.ID, codersdk.UpdateWorkspaceLock{
26652670
Lock: true,
26662671
})
26672672
require.NoError(t, err)
26682673

26692674
workspace = coderdtest.MustWorkspace(t, client, workspace.ID)
26702675
require.NoError(t, err, "fetch provisioned workspace")
2676+
// The template doesn't have a locked_ttl set so this should be nil.
2677+
require.Nil(t, workspace.DeletingAt)
26712678
require.NotNil(t, workspace.LockedAt)
26722679
require.WithinRange(t, *workspace.LockedAt, time.Now().Add(-time.Second*10), time.Now())
2680+
require.Equal(t, lastUsedAt, workspace.LastUsedAt)
26732681

2674-
lastUsedAt := workspace.LastUsedAt
2682+
workspace = coderdtest.MustWorkspace(t, client, workspace.ID)
2683+
lastUsedAt = workspace.LastUsedAt
26752684
err = client.UpdateWorkspaceLock(ctx, workspace.ID, codersdk.UpdateWorkspaceLock{
26762685
Lock: false,
26772686
})
@@ -2680,6 +2689,9 @@ func TestWorkspaceLock(t *testing.T) {
26802689
workspace, err = client.Workspace(ctx, workspace.ID)
26812690
require.NoError(t, err, "fetch provisioned workspace")
26822691
require.Nil(t, workspace.LockedAt)
2692+
// The template doesn't have a locked_ttl set so this should be nil.
2693+
require.Nil(t, workspace.DeletingAt)
2694+
// The last_used_at should get updated when we unlock the workspace.
26832695
require.True(t, workspace.LastUsedAt.After(lastUsedAt))
26842696
})
26852697

‎enterprise/coderd/workspaces_test.go

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -673,3 +673,65 @@ func TestWorkspacesFiltering(t *testing.T) {
673673
assert.Equal(t, workspace.ID, res.Workspaces[0].ID)
674674
})
675675
}
676+
677+
func TestWorkspaceLock(t *testing.T) {
678+
t.Parallel()
679+
680+
t.Run("TemplateLockedTTL", func(t *testing.T) {
681+
t.Parallel()
682+
var (
683+
client = coderdenttest.New(t, &coderdenttest.Options{
684+
Options: &coderdtest.Options{
685+
IncludeProvisionerDaemon: true,
686+
TemplateScheduleStore: &coderd.EnterpriseTemplateScheduleStore{},
687+
},
688+
})
689+
690+
user = coderdtest.CreateFirstUser(t, client)
691+
_ = coderdenttest.AddFullLicense(t, client)
692+
version = coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
693+
_ = coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
694+
lockedTTL = time.Minute
695+
)
696+
697+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID, func(ctr *codersdk.CreateTemplateRequest) {
698+
ctr.LockedTTLMillis = ptr.Ref[int64](lockedTTL.Milliseconds())
699+
})
700+
701+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
702+
_ = coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
703+
704+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
705+
defer cancel()
706+
707+
lastUsedAt := workspace.LastUsedAt
708+
err := client.UpdateWorkspaceLock(ctx, workspace.ID, codersdk.UpdateWorkspaceLock{
709+
Lock: true,
710+
})
711+
require.NoError(t, err)
712+
713+
workspace = coderdtest.MustWorkspace(t, client, workspace.ID)
714+
require.NoError(t, err, "fetch provisioned workspace")
715+
require.NotNil(t, workspace.DeletingAt)
716+
require.NotNil(t, workspace.LockedAt)
717+
require.Equal(t, workspace.LockedAt.Add(lockedTTL), *workspace.DeletingAt)
718+
require.WithinRange(t, *workspace.LockedAt, time.Now().Add(-time.Second*10), time.Now())
719+
// Locking a workspace shouldn't update the last_used_at.
720+
require.Equal(t, lastUsedAt, workspace.LastUsedAt)
721+
722+
workspace = coderdtest.MustWorkspace(t, client, workspace.ID)
723+
lastUsedAt = workspace.LastUsedAt
724+
err = client.UpdateWorkspaceLock(ctx, workspace.ID, codersdk.UpdateWorkspaceLock{
725+
Lock: false,
726+
})
727+
require.NoError(t, err)
728+
729+
workspace, err = client.Workspace(ctx, workspace.ID)
730+
require.NoError(t, err, "fetch provisioned workspace")
731+
require.Nil(t, workspace.LockedAt)
732+
// Unlocking a workspace should cause the deleting_at to be unset.
733+
require.Nil(t, workspace.DeletingAt)
734+
// The last_used_at should get updated when we unlock the workspace.
735+
require.True(t, workspace.LastUsedAt.After(lastUsedAt))
736+
})
737+
}

0 commit comments

Comments
 (0)