shishantbiswas/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-06-02 20:48:20 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
f91475cd51f3cbcb385075fc0add202b5da9a230
coder/enterprise/coderd/schedule/template_test.go
T

1394 lines
49 KiB
Go
Raw Blame History

package schedule_test
import (
"context"
"database/sql"
"encoding/json"
"fmt"
"sync/atomic"
"testing"
"time"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"cdr.dev/slog/v3"
"cdr.dev/slog/v3/sloggers/slogtest"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/dbfake"
"github.com/coder/coder/v2/coderd/database/dbgen"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/notifications"
"github.com/coder/coder/v2/coderd/notifications/notificationstest"
agplschedule "github.com/coder/coder/v2/coderd/schedule"
"github.com/coder/coder/v2/coderd/schedule/cron"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/cryptorand"
"github.com/coder/coder/v2/enterprise/coderd/schedule"
"github.com/coder/coder/v2/provisionersdk/proto"
"github.com/coder/coder/v2/testutil"
"github.com/coder/quartz"
)
func TestTemplateUpdateBuildDeadlines(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
var (
quietUser = dbgen.User(t, db, database.User{
Username: "quiet",
})
noQuietUser = dbgen.User(t, db, database.User{
Username: "no-quiet",
})
file = dbgen.File(t, db, database.File{
CreatedBy: quietUser.ID,
})
templateJob = dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
FileID: file.ID,
InitiatorID: quietUser.ID,
Tags: database.StringMap{
"foo": "bar",
},
})
templateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{
OrganizationID: templateJob.OrganizationID,
CreatedBy: quietUser.ID,
JobID: templateJob.ID,
})
organizationID = templateJob.OrganizationID
)
const userQuietHoursSchedule = "CRON_TZ=UTC 0 0 * * *" // midnight UTC
ctx := testutil.Context(t, testutil.WaitLong)
quietUser, err := db.UpdateUserQuietHoursSchedule(ctx, database.UpdateUserQuietHoursScheduleParams{
ID: quietUser.ID,
QuietHoursSchedule: userQuietHoursSchedule,
})
require.NoError(t, err)
realNow := time.Now().UTC()
nowY, nowM, nowD := realNow.Date()
buildTime := time.Date(nowY, nowM, nowD, 12, 0, 0, 0, time.UTC) // noon today UTC
nextQuietHours := time.Date(nowY, nowM, nowD+1, 0, 0, 0, 0, time.UTC) // midnight tomorrow UTC
defaultTTL := 8 * time.Hour
cases := []struct {
name string
now time.Time
// Before:
deadline time.Time
maxDeadline time.Time
// After:
newDeadline time.Time
newMaxDeadline time.Time
// Config:
noQuietHours bool
// Note that ttl will not influence the new build at all unless it's 0
// AND the build does not have a max deadline post recalculation.
ttl time.Duration
autostopReq *agplschedule.TemplateAutostopRequirement
}{
{
name: "SkippedWorkspaceMaxDeadlineTooSoon",
now: buildTime,
deadline: buildTime,
maxDeadline: buildTime.Add(1 * time.Hour),
// Unchanged since the max deadline is too soon.
newDeadline: buildTime,
newMaxDeadline: buildTime.Add(1 * time.Hour),
ttl: defaultTTL, // no effect
},
{
name: "NewWorkspaceMaxDeadlineBeforeNow",
// After the new max deadline...
now: nextQuietHours.Add(6 * time.Hour),
deadline: buildTime,
// Far into the future...
maxDeadline: nextQuietHours.Add(24 * time.Hour),
newDeadline: buildTime,
// We will use now() + 2 hours if the newly calculated max deadline
// from the workspace build time is before now.
newMaxDeadline: nextQuietHours.Add(8 * time.Hour),
ttl: defaultTTL, // no effect
},
{
name: "NewWorkspaceMaxDeadlineSoon",
// Right before the new max deadline...
now: nextQuietHours.Add(-1 * time.Hour),
deadline: buildTime,
// Far into the future...
maxDeadline: nextQuietHours.Add(24 * time.Hour),
newDeadline: buildTime,
// We will use now() + 2 hours if the newly calculated max deadline
// from the workspace build time is within the next 2 hours.
newMaxDeadline: nextQuietHours.Add(1 * time.Hour),
ttl: defaultTTL, // no effect
},
{
name: "NewWorkspaceMaxDeadlineFuture",
// Well before the new max deadline...
now: nextQuietHours.Add(-6 * time.Hour),
deadline: buildTime,
// Far into the future...
maxDeadline: nextQuietHours.Add(24 * time.Hour),
newDeadline: buildTime,
newMaxDeadline: nextQuietHours,
ttl: defaultTTL, // no effect
},
{
name: "DeadlineAfterNewWorkspaceMaxDeadline",
// Well before the new max deadline...
now: nextQuietHours.Add(-6 * time.Hour),
// Far into the future...
deadline: nextQuietHours.Add(24 * time.Hour),
maxDeadline: nextQuietHours.Add(24 * time.Hour),
// The deadline should match since it is after the new max deadline.
newDeadline: nextQuietHours,
newMaxDeadline: nextQuietHours,
ttl: defaultTTL, // no effect
},
{
// There was a bug if a user has no quiet hours set, and autostop
// req is not turned on, then the max deadline is set to `time.Time{}`.
// This zero value was "in the past", so the workspace deadline would
// be set to "now" + 2 hours.
// This is a mistake because the max deadline being zero means
// there is no max deadline.
name: "MaxDeadlineShouldBeUnset",
now: buildTime,
deadline: buildTime.Add(time.Hour * 8),
maxDeadline: time.Time{}, // No max set
// Should be unchanged
newDeadline: buildTime.Add(time.Hour * 8),
newMaxDeadline: time.Time{},
noQuietHours: true,
autostopReq: &agplschedule.TemplateAutostopRequirement{
DaysOfWeek: 0,
Weeks: 0,
},
ttl: defaultTTL, // no effect
},
{
// A bug existed where MaxDeadline could be set, but deadline was
// `time.Time{}`. This is a logical inconsistency because the "max"
// deadline was ignored.
name: "NoDeadline",
now: buildTime,
deadline: time.Time{},
maxDeadline: time.Time{}, // No max set
// Should be unchanged
newDeadline: time.Time{},
newMaxDeadline: time.Time{},
noQuietHours: true,
autostopReq: &agplschedule.TemplateAutostopRequirement{
DaysOfWeek: 0,
Weeks: 0,
},
ttl: defaultTTL, // no effect
},
{
// Similar to 'NoDeadline' test. This has a MaxDeadline set, so
// the deadline of the workspace should now be set.
name: "WorkspaceDeadlineNowSet",
now: nextQuietHours.Add(-6 * time.Hour),
// Start with unset times
deadline: time.Time{},
maxDeadline: time.Time{},
newDeadline: nextQuietHours,
newMaxDeadline: nextQuietHours,
ttl: defaultTTL, // no effect
},
{
// If the build doesn't have a max_deadline anymore, and there is no
// TTL anymore, then both the deadline and max_deadline should be
// zero.
name: "NoTTLNoDeadlineNoMaxDeadline",
now: buildTime,
deadline: buildTime.Add(time.Hour * 8),
maxDeadline: buildTime.Add(time.Hour * 8),
newDeadline: time.Time{},
newMaxDeadline: time.Time{},
noQuietHours: true,
autostopReq: &agplschedule.TemplateAutostopRequirement{
DaysOfWeek: 0,
Weeks: 0,
},
ttl: 0,
},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
user := quietUser
if c.noQuietHours {
user = noQuietUser
}
t.Log("buildTime", buildTime)
t.Log("nextQuietHours", nextQuietHours)
t.Log("now", c.now)
t.Log("deadline", c.deadline)
t.Log("maxDeadline", c.maxDeadline)
t.Log("newDeadline", c.newDeadline)
t.Log("newMaxDeadline", c.newMaxDeadline)
t.Log("ttl", c.ttl)
template := dbgen.Template(t, db, database.Template{
OrganizationID: organizationID,
ActiveVersionID: templateVersion.ID,
CreatedBy: user.ID,
})
buildResp := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OrganizationID: organizationID,
OwnerID: user.ID,
TemplateID: template.ID,
}).Seed(database.WorkspaceBuild{
TemplateVersionID: templateVersion.ID,
}).ProvisionerState([]byte(must(cryptorand.String(64)))).Succeeded(dbfake.WithJobCompletedAt(buildTime)).Do()
// Assert test invariant: workspace build state must not be empty
var buildProvisionerState []byte
buildProvisionerStateRow, err := db.GetWorkspaceBuildProvisionerStateByID(ctx, buildResp.Build.ID)
require.NoError(t, err)
buildProvisionerState = buildProvisionerStateRow.ProvisionerState
require.NotEmpty(t, buildProvisionerState, "provisioner state must not be empty")
err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
ID: buildResp.Build.ID,
UpdatedAt: buildTime,
Deadline: c.deadline,
MaxDeadline: c.maxDeadline,
})
require.NoError(t, err)
wsBuild, err := db.GetWorkspaceBuildByID(ctx, buildResp.Build.ID)
require.NoError(t, err)
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
userQuietHoursStore, err := schedule.NewEnterpriseUserQuietHoursScheduleStore(userQuietHoursSchedule, true)
require.NoError(t, err)
userQuietHoursStorePtr := &atomic.Pointer[agplschedule.UserQuietHoursScheduleStore]{}
userQuietHoursStorePtr.Store(&userQuietHoursStore)
clock := quartz.NewMock(t)
clock.Set(c.now)
// Set the template policy.
templateScheduleStore := schedule.NewEnterpriseTemplateScheduleStore(userQuietHoursStorePtr, notifications.NewNoopEnqueuer(), logger, clock)
autostopReq := agplschedule.TemplateAutostopRequirement{
// Every day
DaysOfWeek: 0b01111111,
Weeks: 0,
}
if c.autostopReq != nil {
autostopReq = *c.autostopReq
}
_, err = templateScheduleStore.Set(ctx, db, template, agplschedule.TemplateScheduleOptions{
UserAutostartEnabled: false,
UserAutostopEnabled: false,
DefaultTTL: c.ttl,
AutostopRequirement: autostopReq,
FailureTTL: 0,
TimeTilDormant: 0,
TimeTilDormantAutoDelete: 0,
})
require.NoError(t, err)
// Check that the workspace build has the expected deadlines.
newBuild, err := db.GetWorkspaceBuildByID(ctx, wsBuild.ID)
require.NoError(t, err)
require.WithinDuration(t, c.newDeadline, newBuild.Deadline, time.Second, "deadline")
require.WithinDuration(t, c.newMaxDeadline, newBuild.MaxDeadline, time.Second, "max_deadline")
// Check that the new build has the same state as before.
newBuildProvisionerStateRow, err := db.GetWorkspaceBuildProvisionerStateByID(ctx, newBuild.ID)
require.NoError(t, err)
require.Equal(t, buildProvisionerState, newBuildProvisionerStateRow.ProvisionerState, "provisioner state mismatch")
})
}
}
func TestTemplateUpdateBuildDeadlinesSkip(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
var (
user = dbgen.User(t, db, database.User{})
file = dbgen.File(t, db, database.File{
CreatedBy: user.ID,
})
templateJob = dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
FileID: file.ID,
InitiatorID: user.ID,
Tags: database.StringMap{
"foo": "bar",
},
})
templateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{
CreatedBy: user.ID,
JobID: templateJob.ID,
OrganizationID: templateJob.OrganizationID,
})
template = dbgen.Template(t, db, database.Template{
ActiveVersionID: templateVersion.ID,
CreatedBy: user.ID,
OrganizationID: templateJob.OrganizationID,
})
otherTemplate = dbgen.Template(t, db, database.Template{
ActiveVersionID: templateVersion.ID,
CreatedBy: user.ID,
OrganizationID: templateJob.OrganizationID,
})
)
// Create a workspace that will be shared by two builds.
ws := dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: user.ID,
TemplateID: template.ID,
OrganizationID: templateJob.OrganizationID,
})
const userQuietHoursSchedule = "CRON_TZ=UTC 0 0 * * *" // midnight UTC
ctx := testutil.Context(t, testutil.WaitLong)
user, err := db.UpdateUserQuietHoursSchedule(ctx, database.UpdateUserQuietHoursScheduleParams{
ID: user.ID,
QuietHoursSchedule: userQuietHoursSchedule,
})
require.NoError(t, err)
realNow := time.Now().UTC()
nowY, nowM, nowD := realNow.Date()
buildTime := time.Date(nowY, nowM, nowD, 12, 0, 0, 0, time.UTC) // noon today UTC
now := time.Date(nowY, nowM, nowD, 18, 0, 0, 0, time.UTC) // 6pm today UTC
nextQuietHours := time.Date(nowY, nowM, nowD+1, 0, 0, 0, 0, time.UTC) // midnight tomorrow UTC
// A date very far in the future which would definitely be updated.
originalMaxDeadline := time.Date(nowY+1, nowM, nowD, 0, 0, 0, 0, time.UTC)
_ = otherTemplate
builds := []struct {
name string
templateID uuid.UUID
// Nil workspaceID means create a new workspace.
workspaceID uuid.UUID
buildNumber int32
buildStarted bool
buildCompleted bool
buildError bool
shouldBeUpdated bool
// Set below:
wsBuild database.WorkspaceBuild
wsBuildProvisionerState []byte
}{
{
name: "DifferentTemplate",
templateID: otherTemplate.ID,
workspaceID: uuid.Nil,
buildNumber: 1,
buildStarted: true,
buildCompleted: true,
buildError: false,
shouldBeUpdated: false,
},
{
name: "NonStartedBuild",
templateID: template.ID,
workspaceID: uuid.Nil,
buildNumber: 1,
buildStarted: false,
buildCompleted: false,
buildError: false,
shouldBeUpdated: false,
},
{
name: "InProgressBuild",
templateID: template.ID,
workspaceID: uuid.Nil,
buildNumber: 1,
buildStarted: true,
buildCompleted: false,
buildError: false,
shouldBeUpdated: false,
},
{
name: "FailedBuild",
templateID: template.ID,
workspaceID: uuid.Nil,
buildNumber: 1,
buildStarted: true,
buildCompleted: true,
buildError: true,
shouldBeUpdated: false,
},
{
name: "NonLatestBuild",
templateID: template.ID,
workspaceID: ws.ID,
buildNumber: 1,
buildStarted: true,
buildCompleted: true,
buildError: false,
// This build was successful but is not the latest build for this
// workspace, see the next build.
shouldBeUpdated: false,
},
{
name: "LatestBuild",
templateID: template.ID,
workspaceID: ws.ID,
buildNumber: 2,
buildStarted: true,
buildCompleted: true,
buildError: false,
shouldBeUpdated: true,
},
{
name: "LatestBuildOtherWorkspace",
templateID: template.ID,
workspaceID: uuid.Nil,
buildNumber: 1,
buildStarted: true,
buildCompleted: true,
buildError: false,
shouldBeUpdated: true,
},
}
for i, b := range builds {
wsID := b.workspaceID
if wsID == uuid.Nil {
ws := dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: user.ID,
TemplateID: b.templateID,
OrganizationID: templateJob.OrganizationID,
})
wsID = ws.ID
}
job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
FileID: file.ID,
InitiatorID: user.ID,
Provisioner: database.ProvisionerTypeEcho,
Tags: database.StringMap{
wsID.String(): "yeah",
},
OrganizationID: templateJob.OrganizationID,
})
wsBuildProvisionerState := []byte(must(cryptorand.String(64)))
wsBuild := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
WorkspaceID: wsID,
BuildNumber: b.buildNumber,
JobID: job.ID,
InitiatorID: user.ID,
TemplateVersionID: templateVersion.ID,
})
err = db.UpdateWorkspaceBuildProvisionerStateByID(ctx, database.UpdateWorkspaceBuildProvisionerStateByIDParams{
ID: wsBuild.ID,
UpdatedAt: wsBuild.UpdatedAt,
ProvisionerState: wsBuildProvisionerState,
})
require.NoError(t, err)
// Assert test invariant: workspace build state must not be empty
require.NotEmpty(t, wsBuildProvisionerState, "provisioner state must not be empty")
err = db.UpdateWorkspaceBuildDeadlineByID(ctx, database.UpdateWorkspaceBuildDeadlineByIDParams{
ID: wsBuild.ID,
UpdatedAt: buildTime,
Deadline: originalMaxDeadline,
MaxDeadline: originalMaxDeadline,
})
require.NoError(t, err)
wsBuild, err = db.GetWorkspaceBuildByID(ctx, wsBuild.ID)
require.NoError(t, err)
// Assert test invariant: workspace build state must not be empty
require.NotEmpty(t, wsBuildProvisionerState, "provisioner state must not be empty")
builds[i].wsBuildProvisionerState = wsBuildProvisionerState
builds[i].wsBuild = wsBuild
if !b.buildStarted {
continue
}
acquiredJob, err := db.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
OrganizationID: job.OrganizationID,
StartedAt: sql.NullTime{
Time: buildTime,
Valid: true,
},
WorkerID: uuid.NullUUID{
UUID: uuid.New(),
Valid: true,
},
Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
ProvisionerTags: json.RawMessage(fmt.Sprintf(`{%q: "yeah"}`, wsID)),
})
require.NoError(t, err)
require.Equal(t, job.ID, acquiredJob.ID)
if !b.buildCompleted {
continue
}
buildError := ""
if b.buildError {
buildError = "error"
}
err = db.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
ID: job.ID,
CompletedAt: sql.NullTime{
Time: buildTime,
Valid: true,
},
Error: sql.NullString{
String: buildError,
Valid: b.buildError,
},
UpdatedAt: buildTime,
})
require.NoError(t, err)
}
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
userQuietHoursStore, err := schedule.NewEnterpriseUserQuietHoursScheduleStore(userQuietHoursSchedule, true)
require.NoError(t, err)
userQuietHoursStorePtr := &atomic.Pointer[agplschedule.UserQuietHoursScheduleStore]{}
userQuietHoursStorePtr.Store(&userQuietHoursStore)
clock := quartz.NewMock(t)
clock.Set(now)
// Set the template policy.
templateScheduleStore := schedule.NewEnterpriseTemplateScheduleStore(userQuietHoursStorePtr, notifications.NewNoopEnqueuer(), logger, clock)
_, err = templateScheduleStore.Set(ctx, db, template, agplschedule.TemplateScheduleOptions{
UserAutostartEnabled: false,
UserAutostopEnabled: false,
DefaultTTL: 0,
AutostopRequirement: agplschedule.TemplateAutostopRequirement{
// Every day
DaysOfWeek: 0b01111111,
Weeks: 0,
},
FailureTTL: 0,
TimeTilDormant: 0,
TimeTilDormantAutoDelete: 0,
})
require.NoError(t, err)
// Check each build.
for i, b := range builds {
msg := fmt.Sprintf("build %d: %s", i, b.name)
newBuild, err := db.GetWorkspaceBuildByID(ctx, b.wsBuild.ID)
require.NoError(t, err)
if b.shouldBeUpdated {
assert.WithinDuration(t, nextQuietHours, newBuild.Deadline, time.Second, msg)
assert.WithinDuration(t, nextQuietHours, newBuild.MaxDeadline, time.Second, msg)
} else {
assert.WithinDuration(t, originalMaxDeadline, newBuild.Deadline, time.Second, msg)
assert.WithinDuration(t, originalMaxDeadline, newBuild.MaxDeadline, time.Second, msg)
}
newBuildProvisionerStateRow, err := db.GetWorkspaceBuildProvisionerStateByID(ctx, newBuild.ID)
require.NoError(t, err)
assert.Equal(t, builds[i].wsBuildProvisionerState, newBuildProvisionerStateRow.ProvisionerState, "provisioner state mismatch")
}
}
func TestNotifications(t *testing.T) {
t.Parallel()
t.Run("Dormancy", func(t *testing.T) {
t.Parallel()
var (
db, _ = dbtestutil.NewDB(t)
ctx = testutil.Context(t, testutil.WaitLong)
user = dbgen.User(t, db, database.User{})
file = dbgen.File(t, db, database.File{
CreatedBy: user.ID,
})
templateJob = dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
FileID: file.ID,
InitiatorID: user.ID,
Tags: database.StringMap{
"foo": "bar",
},
})
timeTilDormant = time.Minute * 2
templateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{
CreatedBy: user.ID,
JobID: templateJob.ID,
OrganizationID: templateJob.OrganizationID,
})
template = dbgen.Template(t, db, database.Template{
ActiveVersionID: templateVersion.ID,
CreatedBy: user.ID,
OrganizationID: templateJob.OrganizationID,
TimeTilDormant: int64(timeTilDormant),
TimeTilDormantAutoDelete: int64(timeTilDormant),
})
)
// Add two dormant workspaces and one active workspace.
dormantWorkspaces := []database.WorkspaceTable{
dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: user.ID,
TemplateID: template.ID,
OrganizationID: templateJob.OrganizationID,
LastUsedAt: time.Now().Add(-time.Hour),
}),
dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: user.ID,
TemplateID: template.ID,
OrganizationID: templateJob.OrganizationID,
LastUsedAt: time.Now().Add(-time.Hour),
}),
}
dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: user.ID,
TemplateID: template.ID,
OrganizationID: templateJob.OrganizationID,
LastUsedAt: time.Now(),
})
for _, ws := range dormantWorkspaces {
db.UpdateWorkspaceDormantDeletingAt(ctx, database.UpdateWorkspaceDormantDeletingAtParams{
ID: ws.ID,
DormantAt: sql.NullTime{
Time: ws.LastUsedAt.Add(timeTilDormant),
Valid: true,
},
})
}
// Setup dependencies
notifyEnq := notificationstest.FakeEnqueuer{}
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
const userQuietHoursSchedule = "CRON_TZ=UTC 0 0 * * *" // midnight UTC
userQuietHoursStore, err := schedule.NewEnterpriseUserQuietHoursScheduleStore(userQuietHoursSchedule, true)
require.NoError(t, err)
userQuietHoursStorePtr := &atomic.Pointer[agplschedule.UserQuietHoursScheduleStore]{}
userQuietHoursStorePtr.Store(&userQuietHoursStore)
templateScheduleStore := schedule.NewEnterpriseTemplateScheduleStore(userQuietHoursStorePtr, &notifyEnq, logger, nil)
// Lower the dormancy TTL to ensure the schedule recalculates deadlines and
// triggers notifications.
_, err = templateScheduleStore.Set(dbauthz.AsNotifier(ctx), db, template, agplschedule.TemplateScheduleOptions{
TimeTilDormant: timeTilDormant / 2,
TimeTilDormantAutoDelete: timeTilDormant / 2,
})
require.NoError(t, err)
// We should expect a notification for each dormant workspace.
sent := notifyEnq.Sent()
require.Len(t, sent, len(dormantWorkspaces))
for i, dormantWs := range dormantWorkspaces {
require.Equal(t, sent[i].UserID, dormantWs.OwnerID)
require.Equal(t, sent[i].TemplateID, notifications.TemplateWorkspaceMarkedForDeletion)
require.Contains(t, sent[i].Targets, template.ID)
require.Contains(t, sent[i].Targets, dormantWs.ID)
require.Contains(t, sent[i].Targets, dormantWs.OrganizationID)
require.Contains(t, sent[i].Targets, dormantWs.OwnerID)
}
})
// Regression test for https://github.com/coder/coder/issues/20913
// Deleted workspaces should not receive dormancy notifications.
t.Run("DeletedWorkspacesNotNotified", func(t *testing.T) {
t.Parallel()
var (
db, _ = dbtestutil.NewDB(t)
ctx = testutil.Context(t, testutil.WaitLong)
user = dbgen.User(t, db, database.User{})
file = dbgen.File(t, db, database.File{
CreatedBy: user.ID,
})
templateJob = dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
FileID: file.ID,
InitiatorID: user.ID,
Tags: database.StringMap{
"foo": "bar",
},
})
timeTilDormant = time.Minute * 2
templateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{
CreatedBy: user.ID,
JobID: templateJob.ID,
OrganizationID: templateJob.OrganizationID,
})
template = dbgen.Template(t, db, database.Template{
ActiveVersionID: templateVersion.ID,
CreatedBy: user.ID,
OrganizationID: templateJob.OrganizationID,
TimeTilDormant: int64(timeTilDormant),
TimeTilDormantAutoDelete: int64(timeTilDormant),
})
)
// Create a dormant workspace that is NOT deleted.
activeDormantWorkspace := dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: user.ID,
TemplateID: template.ID,
OrganizationID: templateJob.OrganizationID,
LastUsedAt: time.Now().Add(-time.Hour),
})
_, err := db.UpdateWorkspaceDormantDeletingAt(ctx, database.UpdateWorkspaceDormantDeletingAtParams{
ID: activeDormantWorkspace.ID,
DormantAt: sql.NullTime{
Time: activeDormantWorkspace.LastUsedAt.Add(timeTilDormant),
Valid: true,
},
})
require.NoError(t, err)
// Create a dormant workspace that IS deleted.
deletedDormantWorkspace := dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: user.ID,
TemplateID: template.ID,
OrganizationID: templateJob.OrganizationID,
LastUsedAt: time.Now().Add(-time.Hour),
Deleted: true, // Mark as deleted
})
_, err = db.UpdateWorkspaceDormantDeletingAt(ctx, database.UpdateWorkspaceDormantDeletingAtParams{
ID: deletedDormantWorkspace.ID,
DormantAt: sql.NullTime{
Time: deletedDormantWorkspace.LastUsedAt.Add(timeTilDormant),
Valid: true,
},
})
require.NoError(t, err)
// Setup dependencies
notifyEnq := notificationstest.NewFakeEnqueuer()
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
const userQuietHoursSchedule = "CRON_TZ=UTC 0 0 * * *" // midnight UTC
userQuietHoursStore, err := schedule.NewEnterpriseUserQuietHoursScheduleStore(userQuietHoursSchedule, true)
require.NoError(t, err)
userQuietHoursStorePtr := &atomic.Pointer[agplschedule.UserQuietHoursScheduleStore]{}
userQuietHoursStorePtr.Store(&userQuietHoursStore)
templateScheduleStore := schedule.NewEnterpriseTemplateScheduleStore(userQuietHoursStorePtr, notifyEnq, logger, nil)
// Lower the dormancy TTL to ensure the schedule recalculates deadlines and
// triggers notifications.
_, err = templateScheduleStore.Set(dbauthz.AsNotifier(ctx), db, template, agplschedule.TemplateScheduleOptions{
TimeTilDormant: timeTilDormant / 2,
TimeTilDormantAutoDelete: timeTilDormant / 2,
})
require.NoError(t, err)
// We should only receive a notification for the non-deleted dormant workspace.
sent := notifyEnq.Sent()
require.Len(t, sent, 1, "expected exactly 1 notification for the non-deleted workspace")
require.Equal(t, sent[0].UserID, activeDormantWorkspace.OwnerID)
require.Equal(t, sent[0].TemplateID, notifications.TemplateWorkspaceMarkedForDeletion)
require.Contains(t, sent[0].Targets, activeDormantWorkspace.ID)
// Ensure the deleted workspace was NOT notified
for _, notification := range sent {
require.NotContains(t, notification.Targets, deletedDormantWorkspace.ID,
"deleted workspace should not receive notifications")
}
})
// Disabling dormancy auto-deletion should not send "marked for deletion" notifications.
t.Run("DisablingAutoDeleteSendsNoNotifications", func(t *testing.T) {
t.Parallel()
var (
db, _ = dbtestutil.NewDB(t)
user = dbgen.User(t, db, database.User{})
file = dbgen.File(t, db, database.File{
CreatedBy: user.ID,
})
templateJob = dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
FileID: file.ID,
InitiatorID: user.ID,
Tags: database.StringMap{
"foo": "bar",
},
})
timeTilDormant = time.Minute * 2
timeTilDormantAutoDelete = time.Minute * 4
templateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{
CreatedBy: user.ID,
JobID: templateJob.ID,
OrganizationID: templateJob.OrganizationID,
})
template = dbgen.Template(t, db, database.Template{
ActiveVersionID: templateVersion.ID,
CreatedBy: user.ID,
OrganizationID: templateJob.OrganizationID,
})
)
// Given: Dormancy auto deletion is enabled
ctx := testutil.Context(t, testutil.WaitShort)
err := db.UpdateTemplateScheduleByID(ctx, database.UpdateTemplateScheduleByIDParams{
ID: template.ID,
UpdatedAt: dbtime.Now(),
TimeTilDormant: int64(timeTilDormant),
TimeTilDormantAutoDelete: int64(timeTilDormantAutoDelete),
})
require.NoError(t, err)
// Given: A workspace that is marked as dormant
workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: user.ID,
TemplateID: template.ID,
OrganizationID: templateJob.OrganizationID,
LastUsedAt: time.Now().Add(-time.Hour),
})
dormantAt := workspace.LastUsedAt.Add(timeTilDormant)
workspace, err = db.UpdateWorkspaceDormantDeletingAt(ctx, database.UpdateWorkspaceDormantDeletingAtParams{
ID: workspace.ID,
DormantAt: sql.NullTime{
Time: dormantAt,
Valid: true,
},
})
require.NoError(t, err)
require.True(t, workspace.DeletingAt.Valid, "deleting_at should be set when marking workspace dormant")
// Setup dependencies
notifyEnq := notificationstest.NewFakeEnqueuer()
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
const userQuietHoursSchedule = "CRON_TZ=UTC 0 0 * * *" // midnight UTC
userQuietHoursStore, err := schedule.NewEnterpriseUserQuietHoursScheduleStore(userQuietHoursSchedule, true)
require.NoError(t, err)
userQuietHoursStorePtr := &atomic.Pointer[agplschedule.UserQuietHoursScheduleStore]{}
userQuietHoursStorePtr.Store(&userQuietHoursStore)
templateScheduleStore := schedule.NewEnterpriseTemplateScheduleStore(userQuietHoursStorePtr, notifyEnq, logger, nil)
// When: We disable dormancy auto-delete
_, err = templateScheduleStore.Set(dbauthz.AsNotifier(ctx), db, template, agplschedule.TemplateScheduleOptions{
TimeTilDormant: timeTilDormant,
TimeTilDormantAutoDelete: 0,
})
require.NoError(t, err)
// Then: We expect deleting_at to be removed
updated, err := db.GetWorkspaceByID(ctx, workspace.ID)
require.NoError(t, err)
require.False(t, updated.DeletingAt.Valid, "deleting_at should be cleared when auto-deletion is disabled")
// Then: We expect no notifications to have been sent
sent := notifyEnq.Sent()
require.Len(t, sent, 0, "no notifications should be sent when disabling dormancy auto-deletion")
})
}
func TestTemplateTTL(t *testing.T) {
t.Parallel()
tests := []struct {
name string
allowUserAutostop bool
fromTTL time.Duration
toTTL time.Duration
expected sql.NullInt64
}{
{
name: "AllowUserAutostopFalse/ModifyTTLDurationDown",
allowUserAutostop: false,
fromTTL: 24 * time.Hour,
toTTL: 1 * time.Hour,
expected: sql.NullInt64{Valid: true, Int64: int64(1 * time.Hour)},
},
{
name: "AllowUserAutostopFalse/ModifyTTLDurationUp",
allowUserAutostop: false,
fromTTL: 24 * time.Hour,
toTTL: 36 * time.Hour,
expected: sql.NullInt64{Valid: true, Int64: int64(36 * time.Hour)},
},
{
name: "AllowUserAutostopFalse/ModifyTTLDurationSame",
allowUserAutostop: false,
fromTTL: 24 * time.Hour,
toTTL: 24 * time.Hour,
expected: sql.NullInt64{Valid: true, Int64: int64(24 * time.Hour)},
},
{
name: "AllowUserAutostopFalse/DisableTTL",
allowUserAutostop: false,
fromTTL: 24 * time.Hour,
toTTL: 0,
expected: sql.NullInt64{},
},
{
name: "AllowUserAutostopTrue/ModifyTTLDurationDown",
allowUserAutostop: true,
fromTTL: 24 * time.Hour,
toTTL: 1 * time.Hour,
expected: sql.NullInt64{Valid: true, Int64: int64(24 * time.Hour)},
},
{
name: "AllowUserAutostopTrue/ModifyTTLDurationUp",
allowUserAutostop: true,
fromTTL: 24 * time.Hour,
toTTL: 36 * time.Hour,
expected: sql.NullInt64{Valid: true, Int64: int64(24 * time.Hour)},
},
{
name: "AllowUserAutostopTrue/ModifyTTLDurationSame",
allowUserAutostop: true,
fromTTL: 24 * time.Hour,
toTTL: 24 * time.Hour,
expected: sql.NullInt64{Valid: true, Int64: int64(24 * time.Hour)},
},
{
name: "AllowUserAutostopTrue/DisableTTL",
allowUserAutostop: true,
fromTTL: 24 * time.Hour,
toTTL: 0,
expected: sql.NullInt64{Valid: true, Int64: int64(24 * time.Hour)},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
var (
logger = slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
db, _ = dbtestutil.NewDB(t)
ctx = testutil.Context(t, testutil.WaitLong)
user = dbgen.User(t, db, database.User{})
file = dbgen.File(t, db, database.File{CreatedBy: user.ID})
// Create first template
templateJob = dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
FileID: file.ID,
InitiatorID: user.ID,
Tags: database.StringMap{"foo": "bar"},
})
templateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{
CreatedBy: user.ID,
JobID: templateJob.ID,
OrganizationID: templateJob.OrganizationID,
})
template = dbgen.Template(t, db, database.Template{
ActiveVersionID: templateVersion.ID,
CreatedBy: user.ID,
OrganizationID: templateJob.OrganizationID,
AllowUserAutostop: false,
})
// Create second template
otherTTL = tt.fromTTL + 6*time.Hour
otherTemplateJob = dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
FileID: file.ID,
InitiatorID: user.ID,
Tags: database.StringMap{"foo": "bar"},
})
otherTemplateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{
CreatedBy: user.ID,
JobID: otherTemplateJob.ID,
OrganizationID: otherTemplateJob.OrganizationID,
})
otherTemplate = dbgen.Template(t, db, database.Template{
ActiveVersionID: otherTemplateVersion.ID,
CreatedBy: user.ID,
OrganizationID: otherTemplateJob.OrganizationID,
AllowUserAutostop: false,
})
)
// Setup the template schedule store
notifyEnq := notifications.NewNoopEnqueuer()
const userQuietHoursSchedule = "CRON_TZ=UTC 0 0 * * *" // midnight UTC
userQuietHoursStore, err := schedule.NewEnterpriseUserQuietHoursScheduleStore(userQuietHoursSchedule, true)
require.NoError(t, err)
userQuietHoursStorePtr := &atomic.Pointer[agplschedule.UserQuietHoursScheduleStore]{}
userQuietHoursStorePtr.Store(&userQuietHoursStore)
templateScheduleStore := schedule.NewEnterpriseTemplateScheduleStore(userQuietHoursStorePtr, notifyEnq, logger, nil)
// Set both template's default TTL
template, err = templateScheduleStore.Set(ctx, db, template, agplschedule.TemplateScheduleOptions{
DefaultTTL: tt.fromTTL,
})
require.NoError(t, err)
otherTemplate, err = templateScheduleStore.Set(ctx, db, otherTemplate, agplschedule.TemplateScheduleOptions{
DefaultTTL: otherTTL,
})
require.NoError(t, err)
// We create two workspaces here, one with the template we're modifying, the
// other with a different template. We want to ensure we only modify one
// of the workspaces.
workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: user.ID,
TemplateID: template.ID,
OrganizationID: templateJob.OrganizationID,
LastUsedAt: dbtime.Now(),
Ttl: sql.NullInt64{Valid: true, Int64: int64(tt.fromTTL)},
})
otherWorkspace := dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: user.ID,
TemplateID: otherTemplate.ID,
OrganizationID: otherTemplateJob.OrganizationID,
LastUsedAt: dbtime.Now(),
Ttl: sql.NullInt64{Valid: true, Int64: int64(otherTTL)},
})
// Ensure the workspace's start with the correct TTLs
require.Equal(t, sql.NullInt64{Valid: true, Int64: int64(tt.fromTTL)}, workspace.Ttl)
require.Equal(t, sql.NullInt64{Valid: true, Int64: int64(otherTTL)}, otherWorkspace.Ttl)
// Update _only_ the primary template's TTL
_, err = templateScheduleStore.Set(ctx, db, template, agplschedule.TemplateScheduleOptions{
UserAutostopEnabled: tt.allowUserAutostop,
DefaultTTL: tt.toTTL,
})
require.NoError(t, err)
// Verify the primary workspace's TTL is what we expect
ws, err := db.GetWorkspaceByID(ctx, workspace.ID)
require.NoError(t, err)
require.Equal(t, tt.expected, ws.Ttl)
// Verify we haven't changed the other workspace's TTL
ws, err = db.GetWorkspaceByID(ctx, otherWorkspace.ID)
require.NoError(t, err)
require.Equal(t, sql.NullInt64{Valid: true, Int64: int64(otherTTL)}, ws.Ttl)
})
}
t.Run("WorkspaceTTLUpdatedWhenAllowUserAutostopGetsDisabled", func(t *testing.T) {
t.Parallel()
var (
logger = slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
db, _ = dbtestutil.NewDB(t)
ctx = testutil.Context(t, testutil.WaitLong)
user = dbgen.User(t, db, database.User{})
file = dbgen.File(t, db, database.File{CreatedBy: user.ID})
// Create first template
templateJob = dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
FileID: file.ID,
InitiatorID: user.ID,
Tags: database.StringMap{"foo": "bar"},
})
templateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{
CreatedBy: user.ID,
JobID: templateJob.ID,
OrganizationID: templateJob.OrganizationID,
})
template = dbgen.Template(t, db, database.Template{
ActiveVersionID: templateVersion.ID,
CreatedBy: user.ID,
OrganizationID: templateJob.OrganizationID,
})
)
// Setup the template schedule store
notifyEnq := notifications.NewNoopEnqueuer()
const userQuietHoursSchedule = "CRON_TZ=UTC 0 0 * * *" // midnight UTC
userQuietHoursStore, err := schedule.NewEnterpriseUserQuietHoursScheduleStore(userQuietHoursSchedule, true)
require.NoError(t, err)
userQuietHoursStorePtr := &atomic.Pointer[agplschedule.UserQuietHoursScheduleStore]{}
userQuietHoursStorePtr.Store(&userQuietHoursStore)
templateScheduleStore := schedule.NewEnterpriseTemplateScheduleStore(userQuietHoursStorePtr, notifyEnq, logger, nil)
// Enable AllowUserAutostop
template, err = templateScheduleStore.Set(ctx, db, template, agplschedule.TemplateScheduleOptions{
DefaultTTL: 24 * time.Hour,
UserAutostopEnabled: true,
})
require.NoError(t, err)
// Create a workspace with a TTL different than the template's default TTL
workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: user.ID,
TemplateID: template.ID,
OrganizationID: templateJob.OrganizationID,
LastUsedAt: dbtime.Now(),
Ttl: sql.NullInt64{Valid: true, Int64: int64(48 * time.Hour)},
})
// Ensure the workspace start with the correct TTLs
require.Equal(t, sql.NullInt64{Valid: true, Int64: int64(48 * time.Hour)}, workspace.Ttl)
// Disable AllowUserAutostop
template, err = templateScheduleStore.Set(ctx, db, template, agplschedule.TemplateScheduleOptions{
DefaultTTL: 23 * time.Hour,
UserAutostopEnabled: false,
})
require.NoError(t, err)
// Ensure the workspace ends with the correct TTLs
ws, err := db.GetWorkspaceByID(ctx, workspace.ID)
require.NoError(t, err)
require.Equal(t, sql.NullInt64{Valid: true, Int64: int64(23 * time.Hour)}, ws.Ttl)
})
}
func TestTemplateUpdatePrebuilds(t *testing.T) {
t.Parallel()
// Dormant auto-delete configured to 10 hours
dormantAutoDelete := 10 * time.Hour
// TTL configured to 8 hours
ttl := 8 * time.Hour
// Autostop configuration set to everyday at midnight
autostopWeekdays, err := codersdk.WeekdaysToBitmap(codersdk.AllDaysOfWeek)
require.NoError(t, err)
// Autostart configuration set to everyday at midnight
autostartSchedule, err := cron.Weekly("CRON_TZ=UTC 0 0 * * *")
require.NoError(t, err)
autostartWeekdays, err := codersdk.WeekdaysToBitmap(codersdk.AllDaysOfWeek)
require.NoError(t, err)
cases := []struct {
name string
templateSchedule agplschedule.TemplateScheduleOptions
workspaceUpdate func(*testing.T, context.Context, database.Store, time.Time, database.ClaimPrebuiltWorkspaceRow)
assertWorkspace func(*testing.T, context.Context, database.Store, time.Time, bool, database.Workspace)
}{
{
name: "TemplateDormantAutoDeleteUpdatePrebuildAfterClaim",
templateSchedule: agplschedule.TemplateScheduleOptions{
// Template level TimeTilDormantAutodelete set to 10 hours
TimeTilDormantAutoDelete: dormantAutoDelete,
},
workspaceUpdate: func(t *testing.T, ctx context.Context, db database.Store, now time.Time,
workspace database.ClaimPrebuiltWorkspaceRow,
) {
// When: the workspace is marked dormant
dormantWorkspace, err := db.UpdateWorkspaceDormantDeletingAt(ctx, database.UpdateWorkspaceDormantDeletingAtParams{
ID: workspace.ID,
DormantAt: sql.NullTime{
Time: now,
Valid: true,
},
})
require.NoError(t, err)
require.NotNil(t, dormantWorkspace.DormantAt)
},
assertWorkspace: func(t *testing.T, ctx context.Context, db database.Store, now time.Time,
isPrebuild bool, workspace database.Workspace,
) {
if isPrebuild {
// The unclaimed prebuild should have an empty DormantAt and DeletingAt
require.True(t, workspace.DormantAt.Time.IsZero())
require.True(t, workspace.DeletingAt.Time.IsZero())
} else {
// The claimed workspace should have its DormantAt and DeletingAt updated
require.False(t, workspace.DormantAt.Time.IsZero())
require.False(t, workspace.DeletingAt.Time.IsZero())
require.WithinDuration(t, now.UTC(), workspace.DormantAt.Time.UTC(), time.Second)
require.WithinDuration(t, now.Add(dormantAutoDelete).UTC(), workspace.DeletingAt.Time.UTC(), time.Second)
}
},
},
{
name: "TemplateTTLUpdatePrebuildAfterClaim",
templateSchedule: agplschedule.TemplateScheduleOptions{
// Template level TTL can only be set if autostop is disabled for users
DefaultTTL: ttl,
UserAutostopEnabled: false,
},
workspaceUpdate: func(t *testing.T, ctx context.Context, db database.Store, now time.Time,
workspace database.ClaimPrebuiltWorkspaceRow) {
},
assertWorkspace: func(t *testing.T, ctx context.Context, db database.Store, now time.Time,
isPrebuild bool, workspace database.Workspace,
) {
if isPrebuild {
// The unclaimed prebuild should have an empty TTL
require.Equal(t, sql.NullInt64{}, workspace.Ttl)
} else {
// The claimed workspace should have its TTL updated
require.Equal(t, sql.NullInt64{Int64: int64(ttl), Valid: true}, workspace.Ttl)
}
},
},
{
name: "TemplateAutostopUpdatePrebuildAfterClaim",
templateSchedule: agplschedule.TemplateScheduleOptions{
// Template level Autostop set for everyday
AutostopRequirement: agplschedule.TemplateAutostopRequirement{
DaysOfWeek: autostopWeekdays,
Weeks: 0,
},
},
workspaceUpdate: func(t *testing.T, ctx context.Context, db database.Store, now time.Time,
workspace database.ClaimPrebuiltWorkspaceRow) {
},
assertWorkspace: func(t *testing.T, ctx context.Context, db database.Store, now time.Time, isPrebuild bool, workspace database.Workspace) {
if isPrebuild {
// The unclaimed prebuild should have an empty MaxDeadline
prebuildBuild, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, workspace.ID)
require.NoError(t, err)
require.True(t, prebuildBuild.MaxDeadline.IsZero())
} else {
// The claimed workspace should have its MaxDeadline updated
workspaceBuild, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, workspace.ID)
require.NoError(t, err)
require.False(t, workspaceBuild.MaxDeadline.IsZero())
}
},
},
{
name: "TemplateAutostartUpdatePrebuildAfterClaim",
templateSchedule: agplschedule.TemplateScheduleOptions{
// Template level Autostart set for everyday
UserAutostartEnabled: true,
AutostartRequirement: agplschedule.TemplateAutostartRequirement{
DaysOfWeek: autostartWeekdays,
},
},
workspaceUpdate: func(t *testing.T, ctx context.Context, db database.Store, now time.Time, workspace database.ClaimPrebuiltWorkspaceRow) {
// To compute NextStartAt, the workspace must have a valid autostart schedule
err = db.UpdateWorkspaceAutostart(ctx, database.UpdateWorkspaceAutostartParams{
ID: workspace.ID,
AutostartSchedule: sql.NullString{
String: autostartSchedule.String(),
Valid: true,
},
})
require.NoError(t, err)
},
assertWorkspace: func(t *testing.T, ctx context.Context, db database.Store, now time.Time, isPrebuild bool, workspace database.Workspace) {
if isPrebuild {
// The unclaimed prebuild should have an empty NextStartAt
require.True(t, workspace.NextStartAt.Time.IsZero())
} else {
// The claimed workspace should have its NextStartAt updated
require.False(t, workspace.NextStartAt.Time.IsZero())
}
},
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
clock := quartz.NewMock(t)
clock.Set(dbtime.Now())
// Setup
var (
logger = slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
db, _ = dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure())
ctx = testutil.Context(t, testutil.WaitLong)
user = dbgen.User(t, db, database.User{})
)
// Setup the template schedule store
notifyEnq := notifications.NewNoopEnqueuer()
const userQuietHoursSchedule = "CRON_TZ=UTC 0 0 * * *" // midnight UTC
userQuietHoursStore, err := schedule.NewEnterpriseUserQuietHoursScheduleStore(userQuietHoursSchedule, true)
require.NoError(t, err)
userQuietHoursStorePtr := &atomic.Pointer[agplschedule.UserQuietHoursScheduleStore]{}
userQuietHoursStorePtr.Store(&userQuietHoursStore)
templateScheduleStore := schedule.NewEnterpriseTemplateScheduleStore(userQuietHoursStorePtr, notifyEnq, logger, clock)
// Given: a template and a template version with preset and a prebuilt workspace
presetID := uuid.New()
org := dbfake.Organization(t, db).Do()
tv := dbfake.TemplateVersion(t, db).Seed(database.TemplateVersion{
OrganizationID: org.Org.ID,
CreatedBy: user.ID,
}).Preset(database.TemplateVersionPreset{
ID: presetID,
DesiredInstances: sql.NullInt32{
Int32: 1,
Valid: true,
},
}).Do()
workspaceBuild := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
OwnerID: database.PrebuildsSystemUserID,
TemplateID: tv.Template.ID,
OrganizationID: tv.Template.OrganizationID,
}).Seed(database.WorkspaceBuild{
TemplateVersionID: tv.TemplateVersion.ID,
TemplateVersionPresetID: uuid.NullUUID{
UUID: presetID,
Valid: true,
},
}).WithAgent(func(agent []*proto.Agent) []*proto.Agent {
return agent
}).Do()
// Mark the prebuilt workspace's agent as ready so the prebuild can be claimed
agentCtx := testutil.Context(t, testutil.WaitLong)
agent, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(agentCtx, uuid.MustParse(workspaceBuild.AgentToken))
require.NoError(t, err)
err = db.UpdateWorkspaceAgentLifecycleStateByID(agentCtx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{
ID: agent.WorkspaceAgent.ID,
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
})
require.NoError(t, err)
// Given: a prebuilt workspace
prebuild, err := db.GetWorkspaceByID(ctx, workspaceBuild.Workspace.ID)
require.NoError(t, err)
tc.assertWorkspace(t, ctx, db, clock.Now(), true, prebuild)
// When: the template schedule is updated
_, err = templateScheduleStore.Set(ctx, db, tv.Template, tc.templateSchedule)
require.NoError(t, err)
// Then: lifecycle parameters must remain unset while the prebuild is unclaimed
prebuild, err = db.GetWorkspaceByID(ctx, workspaceBuild.Workspace.ID)
require.NoError(t, err)
tc.assertWorkspace(t, ctx, db, clock.Now(), true, prebuild)
// Given: the prebuilt workspace is claimed by a user
claimedWorkspace := dbgen.ClaimPrebuild(
t, db,
clock.Now(),
user.ID,
"claimedWorkspace-autostop",
presetID,
sql.NullString{},
sql.NullTime{},
sql.NullInt64{})
require.Equal(t, prebuild.ID, claimedWorkspace.ID)
// Given: the workspace level configurations are properly set in order to ensure the
// lifecycle parameters are updated
tc.workspaceUpdate(t, ctx, db, clock.Now(), claimedWorkspace)
// When: the template schedule is updated
_, err = templateScheduleStore.Set(ctx, db, tv.Template, tc.templateSchedule)
require.NoError(t, err)
// Then: the workspace should have its lifecycle parameters updated
workspace, err := db.GetWorkspaceByID(ctx, claimedWorkspace.ID)
require.NoError(t, err)
tc.assertWorkspace(t, ctx, db, clock.Now(), false, workspace)
})
}
}
func must[V any](v V, err error) V {
if err != nil {
panic(err)
}
return v
}
Reference in New Issue View Git Blame Copy Permalink