chore: add test for activitybump max_ttl

coder · deansheather · Mar 7, 2023 · Feb 8, 2023 · Feb 15, 2023 · Feb 20, 2023
commit 25f7d2c5a6671fc7e9ce8fa5e7daf0b9da1565d1
diff --git a/cli/testdata/coder_list_--output_json.golden b/cli/testdata/coder_list_--output_json.golden
@@ -39,6 +39,7 @@
       "reason": "initiator",
       "resources": [],
       "deadline": "[timestamp]",
+      "max_deadline": null,
       "status": "running",
       "daily_cost": 0
     },

diff --git a/coderd/activitybump_test.go b/coderd/activitybump_test.go
@@ -5,34 +5,62 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/coder/coderd/coderdtest"
 	"github.com/coder/coder/coderd/database"
+	"github.com/coder/coder/coderd/provisionerdserver"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/testutil"
 )
 
+type mockTemplateScheduleStore struct {
+	getFn func(ctx context.Context, db database.Store, templateID uuid.UUID) (provisionerdserver.TemplateScheduleOptions, error)
+}
+
+var _ provisionerdserver.TemplateScheduleStore = mockTemplateScheduleStore{}
+
+func (m mockTemplateScheduleStore) GetTemplateScheduleOptions(ctx context.Context, db database.Store, templateID uuid.UUID) (provisionerdserver.TemplateScheduleOptions, error) {
+	return m.getFn(ctx, db, templateID)
+}
+
+func (mockTemplateScheduleStore) SetTemplateScheduleOptions(ctx context.Context, db database.Store, template database.Template, options provisionerdserver.TemplateScheduleOptions) (database.Template, error) {
+	return provisionerdserver.NewAGPLTemplateScheduleStore().SetTemplateScheduleOptions(ctx, db, template, options)
+}
+
 func TestWorkspaceActivityBump(t *testing.T) {
 	t.Parallel()
 
 	ctx := context.Background()
 
-	const ttl = time.Minute
-
-	setupActivityTest := func(t *testing.T) (client *codersdk.Client, workspace codersdk.Workspace, assertBumped func(want bool)) {
-		ttlMillis := int64(ttl / time.Millisecond)
+	setupActivityTest := func(t *testing.T, maxDeadline ...time.Duration) (client *codersdk.Client, workspace codersdk.Workspace, assertBumped func(want bool)) {
+		const ttl = time.Minute
+		maxTTL := time.Duration(0)
+		if len(maxDeadline) > 0 {
+			maxTTL = maxDeadline[0]
+		}
 
 		client = coderdtest.New(t, &coderdtest.Options{
 			AppHostname:              proxyTestSubdomainRaw,
 			IncludeProvisionerDaemon: true,
 			// Agent stats trigger the activity bump, so we want to report
 			// very frequently in tests.
 			AgentStatsRefreshInterval: time.Millisecond * 100,
+			TemplateScheduleStore: mockTemplateScheduleStore{
+				getFn: func(ctx context.Context, db database.Store, templateID uuid.UUID) (provisionerdserver.TemplateScheduleOptions, error) {
+					return provisionerdserver.TemplateScheduleOptions{
+						UserSchedulingEnabled: true,
+						DefaultTTL:            ttl,
+						MaxTTL:                maxTTL,
+					}, nil
+				},
+			},
 		})
 		user := coderdtest.CreateFirstUser(t, client)
 
+		ttlMillis := int64(ttl / time.Millisecond)
 		workspace = createWorkspaceWithApps(t, client, user.OrganizationID, "", 1234, func(cwr *codersdk.CreateWorkspaceRequest) {
 			cwr.TTLMillis = &ttlMillis
 		})
@@ -42,10 +70,21 @@ func TestWorkspaceActivityBump(t *testing.T) {
 		require.NoError(t, err)
 		require.WithinDuration(t,
 			time.Now().Add(time.Duration(ttlMillis)*time.Millisecond),
-			workspace.LatestBuild.Deadline.Time, testutil.WaitMedium,
+			workspace.LatestBuild.Deadline.Time,
+			testutil.WaitMedium,
 		)
 		firstDeadline := workspace.LatestBuild.Deadline.Time
 
+		if maxTTL != 0 {
+			require.WithinDuration(t,
+				time.Now().Add(maxTTL),
+				workspace.LatestBuild.MaxDeadline.Time,
+				testutil.WaitMedium,
+			)
+		} else {
+			require.True(t, workspace.LatestBuild.MaxDeadline.Time.IsZero())
+		}
+
 		_ = coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
 
 		return client, workspace, func(want bool) {
@@ -74,6 +113,12 @@ func TestWorkspaceActivityBump(t *testing.T) {
 				"deadline %v never updated", firstDeadline,
 			)
 
+			// If the workspace has a max deadline, the deadline must not exceed
+			// it.
+			if maxTTL != 0 && database.Now().Add(ttl).After(workspace.LatestBuild.MaxDeadline.Time) {
+				require.Equal(t, workspace.LatestBuild.Deadline.Time, workspace.LatestBuild.MaxDeadline.Time)
+				return
+			}
 			require.WithinDuration(t, database.Now().Add(ttl), workspace.LatestBuild.Deadline.Time, 3*time.Second)
 		}
 	}
@@ -111,4 +156,34 @@ func TestWorkspaceActivityBump(t *testing.T) {
 
 		assertBumped(false)
 	})
+
+	t.Run("NotExceedMaxDeadline", func(t *testing.T) {
+		t.Parallel()
+
+		// Set the max deadline to be in 61 seconds. We bump by 1 minute, so we
+		// should expect the deadline to match the max deadline exactly.
+		client, workspace, assertBumped := setupActivityTest(t, 61*time.Second)
+
+		// Bump by dialing the workspace and sending traffic.
+		resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+		conn, err := client.DialWorkspaceAgent(ctx, resources[0].Agents[0].ID, &codersdk.DialWorkspaceAgentOptions{
+			Logger: slogtest.Make(t, nil),
+		})
+		require.NoError(t, err)
+		defer conn.Close()
+
+		// Must send network traffic after a few seconds to surpass bump threshold.
+		time.Sleep(time.Second * 3)
+		sshConn, err := conn.SSHClient(ctx)
+		require.NoError(t, err)
+		_ = sshConn.Close()
+
+		assertBumped(true)
+
+		// Double check that the workspace build's deadline is equal to the
+		// max deadline.
+		workspace, err = client.Workspace(ctx, workspace.ID)
+		require.NoError(t, err)
+		require.Equal(t, workspace.LatestBuild.Deadline.Time, workspace.LatestBuild.MaxDeadline.Time)
+	})
 }
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -65,6 +65,7 @@ import (
 	"github.com/coder/coder/coderd/gitsshkey"
 	"github.com/coder/coder/coderd/httpapi"
 	"github.com/coder/coder/coderd/httpmw"
+	"github.com/coder/coder/coderd/provisionerdserver"
 	"github.com/coder/coder/coderd/rbac"
 	"github.com/coder/coder/coderd/telemetry"
 	"github.com/coder/coder/coderd/updatecheck"
@@ -85,22 +86,23 @@ type Options struct {
 	// AccessURL denotes a custom access URL. By default we use the httptest
 	// server's URL. Setting this may result in unexpected behavior (especially
 	// with running agents).
-	AccessURL            *url.URL
-	AppHostname          string
-	AWSCertificates      awsidentity.Certificates
-	Authorizer           rbac.Authorizer
-	AzureCertificates    x509.VerifyOptions
-	GithubOAuth2Config   *coderd.GithubOAuth2Config
-	RealIPConfig         *httpmw.RealIPConfig
-	OIDCConfig           *coderd.OIDCConfig
-	GoogleTokenValidator *idtoken.Validator
-	SSHKeygenAlgorithm   gitsshkey.Algorithm
-	AutobuildTicker      <-chan time.Time
-	AutobuildStats       chan<- executor.Stats
-	Auditor              audit.Auditor
-	TLSCertificates      []tls.Certificate
-	GitAuthConfigs       []*gitauth.Config
-	TrialGenerator       func(context.Context, string) error
+	AccessURL             *url.URL
+	AppHostname           string
+	AWSCertificates       awsidentity.Certificates
+	Authorizer            rbac.Authorizer
+	AzureCertificates     x509.VerifyOptions
+	GithubOAuth2Config    *coderd.GithubOAuth2Config
+	RealIPConfig          *httpmw.RealIPConfig
+	OIDCConfig            *coderd.OIDCConfig
+	GoogleTokenValidator  *idtoken.Validator
+	SSHKeygenAlgorithm    gitsshkey.Algorithm
+	AutobuildTicker       <-chan time.Time
+	AutobuildStats        chan<- executor.Stats
+	Auditor               audit.Auditor
+	TLSCertificates       []tls.Certificate
+	GitAuthConfigs        []*gitauth.Config
+	TrialGenerator        func(context.Context, string) error
+	TemplateScheduleStore provisionerdserver.TemplateScheduleStore
 
 	// All rate limits default to -1 (unlimited) in tests if not set.
 	APIRateLimit   int
@@ -287,22 +289,23 @@ func NewOptions(t *testing.T, options *Options) (func(http.Handler), context.Can
 			Pubsub:                         options.Pubsub,
 			GitAuthConfigs:                 options.GitAuthConfigs,
 
-			Auditor:              options.Auditor,
-			AWSCertificates:      options.AWSCertificates,
-			AzureCertificates:    options.AzureCertificates,
-			GithubOAuth2Config:   options.GithubOAuth2Config,
-			RealIPConfig:         options.RealIPConfig,
-			OIDCConfig:           options.OIDCConfig,
-			GoogleTokenValidator: options.GoogleTokenValidator,
-			SSHKeygenAlgorithm:   options.SSHKeygenAlgorithm,
-			DERPServer:           derpServer,
-			APIRateLimit:         options.APIRateLimit,
-			LoginRateLimit:       options.LoginRateLimit,
-			FilesRateLimit:       options.FilesRateLimit,
-			Authorizer:           options.Authorizer,
-			Telemetry:            telemetry.NewNoop(),
-			TLSCertificates:      options.TLSCertificates,
-			TrialGenerator:       options.TrialGenerator,
+			Auditor:               options.Auditor,
+			AWSCertificates:       options.AWSCertificates,
+			AzureCertificates:     options.AzureCertificates,
+			GithubOAuth2Config:    options.GithubOAuth2Config,
+			RealIPConfig:          options.RealIPConfig,
+			OIDCConfig:            options.OIDCConfig,
+			GoogleTokenValidator:  options.GoogleTokenValidator,
+			SSHKeygenAlgorithm:    options.SSHKeygenAlgorithm,
+			DERPServer:            derpServer,
+			APIRateLimit:          options.APIRateLimit,
+			LoginRateLimit:        options.LoginRateLimit,
+			FilesRateLimit:        options.FilesRateLimit,
+			Authorizer:            options.Authorizer,
+			Telemetry:             telemetry.NewNoop(),
+			TemplateScheduleStore: options.TemplateScheduleStore,
+			TLSCertificates:       options.TLSCertificates,
+			TrialGenerator:        options.TrialGenerator,
 			DERPMap: &tailcfg.DERPMap{
 				Regions: map[int]*tailcfg.DERPRegion{
 					1: {

diff --git a/coderd/workspacebuilds.go b/coderd/workspacebuilds.go
@@ -1139,6 +1139,7 @@ func (api *API) convertWorkspaceBuild(
 		InitiatorUsername:   initiator.Username,
 		Job:                 apiJob,
 		Deadline:            codersdk.NewNullTime(build.Deadline, !build.Deadline.IsZero()),
+		MaxDeadline:         codersdk.NewNullTime(build.MaxDeadline, !build.MaxDeadline.IsZero()),
 		Reason:              codersdk.BuildReason(build.Reason),
 		Resources:           apiResources,
 		Status:              convertWorkspaceStatus(apiJob.Status, transition),

diff --git a/codersdk/workspacebuilds.go b/codersdk/workspacebuilds.go
@@ -68,6 +68,7 @@ type WorkspaceBuild struct {
 	Reason              BuildReason         `db:"reason" json:"reason" enums:"initiator,autostart,autostop"`
 	Resources           []WorkspaceResource `json:"resources"`
 	Deadline            NullTime            `json:"deadline,omitempty" format:"date-time"`
+	MaxDeadline         NullTime            `json:"max_deadline,omitempty" format:"date-time"`
 	Status              WorkspaceStatus     `json:"status" enums:"pending,starting,running,stopping,stopped,failed,canceling,canceled,deleting,deleted"`
 	DailyCost           int32               `json:"daily_cost"`
 }