chore: add agentapi tests (#11269)

coderd/agentapi/stats.go

@@ -9,57 +9,71 @@ import (
 	"golang.org/x/xerrors"
 	"google.golang.org/protobuf/types/known/durationpb"
 
+	"github.com/google/uuid"
+
 	"cdr.dev/slog"
 	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/autobuild"
-	"github.com/coder/coder/v2/coderd/batchstats"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/prometheusmetrics"
 	"github.com/coder/coder/v2/coderd/schedule"
 )
 
+type StatsBatcher interface {
+	Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error
+}
+
 type StatsAPI struct {
 	AgentFn                   func(context.Context) (database.WorkspaceAgent, error)
 	Database                  database.Store
 	Log                       slog.Logger
-	StatsBatcher              *batchstats.Batcher
+	StatsBatcher              StatsBatcher
 	TemplateScheduleStore     *atomic.Pointer[schedule.TemplateScheduleStore]
 	AgentStatsRefreshInterval time.Duration
 	UpdateAgentMetricsFn      func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
+
+	TimeNowFn func() time.Time // defaults to dbtime.Now()
+}
+
+func (a *StatsAPI) now() time.Time {
+	if a.TimeNowFn != nil {
+		return a.TimeNowFn()
+	}
+	return dbtime.Now()
 }
 
 func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsRequest) (*agentproto.UpdateStatsResponse, error) {
+	// An empty stat means it's just looking for the report interval.
+	res := &agentproto.UpdateStatsResponse{
+		ReportInterval: durationpb.New(a.AgentStatsRefreshInterval),
+	}
+	if req.Stats == nil || len(req.Stats.ConnectionsByProto) == 0 {
+		return res, nil
+	}
+
 	workspaceAgent, err := a.AgentFn(ctx)
 	if err != nil {
 		return nil, err
 	}
-	row, err := a.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID)
+	getWorkspaceAgentByIDRow, err := a.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID)
 	if err != nil {
 		return nil, xerrors.Errorf("get workspace by agent ID %q: %w", workspaceAgent.ID, err)
 	}
-	workspace := row.Workspace
-
-	res := &agentproto.UpdateStatsResponse{
-		ReportInterval: durationpb.New(a.AgentStatsRefreshInterval),
-	}
-
-	// An empty stat means it's just looking for the report interval.
-	if len(req.Stats.ConnectionsByProto) == 0 {
-		return res, nil
-	}
-
+	workspace := getWorkspaceAgentByIDRow.Workspace
 	a.Log.Debug(ctx, "read stats report",
 		slog.F("interval", a.AgentStatsRefreshInterval),
 		slog.F("workspace_id", workspace.ID),
 		slog.F("payload", req),
 	)
 
+	now := a.now()
 	if req.Stats.ConnectionCount > 0 {
 		var nextAutostart time.Time
 		if workspace.AutostartSchedule.String != "" {
 			templateSchedule, err := (*(a.TemplateScheduleStore.Load())).Get(ctx, a.Database, workspace.TemplateID)
-			// If the template schedule fails to load, just default to bumping without the next trasition and log it.
+			// If the template schedule fails to load, just default to bumping
+			// without the next transition and log it.
 			if err != nil {
 				a.Log.Error(ctx, "failed to load template schedule bumping activity, defaulting to bumping by 60min",
 					slog.F("workspace_id", workspace.ID),
@@ -67,7 +81,7 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
 					slog.Error(err),
 				)
 			} else {
-				next, allowed := autobuild.NextAutostartSchedule(time.Now(), workspace.AutostartSchedule.String, templateSchedule)
+				next, allowed := autobuild.NextAutostartSchedule(now, workspace.AutostartSchedule.String, templateSchedule)
 				if allowed {
 					nextAutostart = next
 				}
@@ -76,13 +90,12 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
 		ActivityBumpWorkspace(ctx, a.Log.Named("activity_bump"), a.Database, workspace.ID, nextAutostart)
 	}
 
-	now := dbtime.Now()
-
 	var errGroup errgroup.Group
 	errGroup.Go(func() error {
-		if err := a.StatsBatcher.Add(time.Now(), workspaceAgent.ID, workspace.TemplateID, workspace.OwnerID, workspace.ID, req.Stats); err != nil {
-			a.Log.Error(ctx, "failed to add stats to batcher", slog.Error(err))
-			return xerrors.Errorf("can't insert workspace agent stat: %w", err)
+		err := a.StatsBatcher.Add(now, workspaceAgent.ID, workspace.TemplateID, workspace.OwnerID, workspace.ID, req.Stats)
+		if err != nil {
+			a.Log.Error(ctx, "add agent stats to batcher", slog.Error(err))
+			return xerrors.Errorf("insert workspace agent stats batch: %w", err)
 		}
 		return nil
 	})
@@ -92,7 +105,7 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
 			LastUsedAt: now,
 		})
 		if err != nil {
-			return xerrors.Errorf("can't update workspace LastUsedAt: %w", err)
+			return xerrors.Errorf("update workspace LastUsedAt: %w", err)
 		}
 		return nil
 	})
@@ -100,14 +113,14 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
 		errGroup.Go(func() error {
 			user, err := a.Database.GetUserByID(ctx, workspace.OwnerID)
 			if err != nil {
-				return xerrors.Errorf("can't get user: %w", err)
+				return xerrors.Errorf("get user: %w", err)
 			}
 
 			a.UpdateAgentMetricsFn(ctx, prometheusmetrics.AgentMetricLabels{
 				Username:      user.Username,
 				WorkspaceName: workspace.Name,
 				AgentName:     workspaceAgent.Name,
-				TemplateName:  row.TemplateName,
+				TemplateName:  getWorkspaceAgentByIDRow.TemplateName,
 			}, req.Stats.Metrics)
 			return nil
 		})