coderd/metricscache/metricscache_test.go
Rama main · réplica de lectura del repositorio original.
package metricscache_test
import (
"context"
"database/sql"
"sync/atomic"
"testing"
"time"
"github.com/google/uuid"
"github.com/prometheus/client_golang/prometheus"
"github.com/stretchr/testify/require"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"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/metricscache"
"github.com/coder/coder/v2/coderd/rbac"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/testutil"
"github.com/coder/quartz"
)
func date(year, month, day int) time.Time {
return time.Date(year, time.Month(month), day, 0, 0, 0, 0, time.UTC)
}
func newMetricsCache(t *testing.T, log slog.Logger, clock quartz.Clock, intervals metricscache.Intervals, usage bool) (*metricscache.Cache, database.Store) {
t.Helper()
accessControlStore := &atomic.Pointer[dbauthz.AccessControlStore]{}
var acs dbauthz.AccessControlStore = dbauthz.AGPLTemplateAccessControlStore{}
accessControlStore.Store(&acs)
var (
auth = rbac.NewStrictCachingAuthorizer(prometheus.NewRegistry())
db, _ = dbtestutil.NewDB(t)
dbauth = dbauthz.New(db, auth, log, accessControlStore)
cache = metricscache.New(dbauth, log, clock, intervals, usage)
)
t.Cleanup(func() { cache.Close() })
return cache, db
}
func TestCache_TemplateWorkspaceOwners(t *testing.T) {
t.Parallel()
var (
ctx = testutil.Context(t, testutil.WaitShort)
log = testutil.Logger(t)
clock = quartz.NewMock(t)
)
trapTickerFunc := clock.Trap().TickerFunc("metricscache")
defer trapTickerFunc.Close()
cache, db := newMetricsCache(t, log, clock, metricscache.Intervals{
TemplateBuildTimes: time.Minute,
}, false)
org := dbgen.Organization(t, db, database.Organization{})
user1 := dbgen.User(t, db, database.User{})
user2 := dbgen.User(t, db, database.User{})
template := dbgen.Template(t, db, database.Template{
OrganizationID: org.ID,
Provisioner: database.ProvisionerTypeEcho,
CreatedBy: user1.ID,
})
// Wait for both ticker functions to be created (template build times and deployment stats)
trapTickerFunc.MustWait(ctx).MustRelease(ctx)
trapTickerFunc.MustWait(ctx).MustRelease(ctx)
clock.Advance(time.Minute).MustWait(ctx)
count, ok := cache.TemplateWorkspaceOwners(template.ID)
require.True(t, ok, "TemplateWorkspaceOwners should be populated")
require.Equal(t, 0, count, "should have 0 owners initially")
dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
TemplateID: template.ID,
OwnerID: user1.ID,
})
clock.Advance(time.Minute).MustWait(ctx)
count, _ = cache.TemplateWorkspaceOwners(template.ID)
require.Equal(t, 1, count, "should have 1 owner after adding workspace")
workspace2 := dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
TemplateID: template.ID,
OwnerID: user2.ID,
})
clock.Advance(time.Minute).MustWait(ctx)
count, _ = cache.TemplateWorkspaceOwners(template.ID)
require.Equal(t, 2, count, "should have 2 owners after adding second workspace")
// 3rd workspace should not be counted since we have the same owner as workspace2.
dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
TemplateID: template.ID,
OwnerID: user1.ID,
})
db.UpdateWorkspaceDeletedByID(context.Background(), database.UpdateWorkspaceDeletedByIDParams{
ID: workspace2.ID,
Deleted: true,
})
clock.Advance(time.Minute).MustWait(ctx)
count, _ = cache.TemplateWorkspaceOwners(template.ID)
require.Equal(t, 1, count, "should have 1 owner after deleting workspace")
}
func clockTime(t time.Time, hour, minute, sec int) time.Time {
return time.Date(t.Year(), t.Month(), t.Day(), hour, minute, sec, t.Nanosecond(), t.Location())
}
func requireBuildTimeStatsEmpty(t *testing.T, stats codersdk.TemplateBuildTimeStats) {
require.Empty(t, stats[codersdk.WorkspaceTransitionStart])
require.Empty(t, stats[codersdk.WorkspaceTransitionStop])
require.Empty(t, stats[codersdk.WorkspaceTransitionDelete])
}
func TestCache_BuildTime(t *testing.T) {
t.Parallel()
someDay := date(2022, 10, 1)
type jobParams struct {
startedAt time.Time
completedAt time.Time
}
type args struct {
rows []jobParams
transition database.WorkspaceTransition
}
type want struct {
buildTimeMs int64
loads bool
}
tests := []struct {
name string
args args
want want
}{
{"empty", args{}, want{-1, false}},
{
"one/start", args{
rows: []jobParams{
{
startedAt: clockTime(someDay, 10, 1, 0),
completedAt: clockTime(someDay, 10, 1, 10),
},
},
transition: database.WorkspaceTransitionStart,
}, want{10 * 1000, true},
},
{
"two/stop", args{
rows: []jobParams{
{
startedAt: clockTime(someDay, 10, 1, 0),
completedAt: clockTime(someDay, 10, 1, 10),
},
{
startedAt: clockTime(someDay, 10, 1, 0),
completedAt: clockTime(someDay, 10, 1, 50),
},
},
transition: database.WorkspaceTransitionStop,
}, want{10 * 1000, true},
},
{
"three/delete", args{
rows: []jobParams{
{
startedAt: clockTime(someDay, 10, 1, 0),
completedAt: clockTime(someDay, 10, 1, 10),
},
{
startedAt: clockTime(someDay, 10, 1, 0),
completedAt: clockTime(someDay, 10, 1, 50),
},
{
startedAt: clockTime(someDay, 10, 1, 0),
completedAt: clockTime(someDay, 10, 1, 20),
},
},
transition: database.WorkspaceTransitionDelete,
}, want{20 * 1000, true},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
var (
ctx = testutil.Context(t, testutil.WaitShort)
log = testutil.Logger(t)
clock = quartz.NewMock(t)
)
clock.Set(someDay)
trapTickerFunc := clock.Trap().TickerFunc("metricscache")
defer trapTickerFunc.Close()
cache, db := newMetricsCache(t, log, clock, metricscache.Intervals{
TemplateBuildTimes: time.Minute,
}, false)
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
template := dbgen.Template(t, db, database.Template{
CreatedBy: user.ID,
OrganizationID: org.ID,
})
templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
OrganizationID: org.ID,
CreatedBy: user.ID,
TemplateID: uuid.NullUUID{UUID: template.ID, Valid: true},
})
workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
OwnerID: user.ID,
TemplateID: template.ID,
})
gotStats := cache.TemplateBuildTimeStats(template.ID)
requireBuildTimeStatsEmpty(t, gotStats)
for buildNumber, row := range tt.args.rows {
job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
OrganizationID: org.ID,
InitiatorID: user.ID,
Type: database.ProvisionerJobTypeWorkspaceBuild,
StartedAt: sql.NullTime{Time: row.startedAt, Valid: true},
CompletedAt: sql.NullTime{Time: row.completedAt, Valid: true},
})
dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
BuildNumber: int32(1 + buildNumber), // nolint:gosec
WorkspaceID: workspace.ID,
InitiatorID: user.ID,
TemplateVersionID: templateVersion.ID,
JobID: job.ID,
Transition: tt.args.transition,
})
}
// Wait for both ticker functions to be created (template build times and deployment stats)
trapTickerFunc.MustWait(ctx).MustRelease(ctx)
trapTickerFunc.MustWait(ctx).MustRelease(ctx)
clock.Advance(time.Minute).MustWait(ctx)
if tt.want.loads {
wantTransition := codersdk.WorkspaceTransition(tt.args.transition)
gotStats := cache.TemplateBuildTimeStats(template.ID)
ts := gotStats[wantTransition]
require.NotNil(t, ts.P50, "P50 should be set for %v", wantTransition)
require.Equal(t, tt.want.buildTimeMs, *ts.P50, "P50 should match expected value for %v", wantTransition)
for transition, ts := range gotStats {
if transition == wantTransition {
// Checked above
continue
}
require.Empty(t, ts, "%v", transition)
}
} else {
stats := cache.TemplateBuildTimeStats(template.ID)
requireBuildTimeStatsEmpty(t, stats)
}
})
}
}
func TestCache_DeploymentStats(t *testing.T) {
t.Parallel()
for name, usage := range map[string]bool{"Stats": false, "Usage": true} {
t.Run(name, func(t *testing.T) {
t.Parallel()
var (
ctx = testutil.Context(t, testutil.WaitLong)
clock = quartz.NewMock(t)
now = dbtime.Now().Truncate(time.Minute)
tickerTrap = clock.Trap().TickerFunc("metricscache")
)
clock.Set(now)
defer tickerTrap.Close()
cache, db := newMetricsCache(t, testutil.Logger(t), clock, metricscache.Intervals{
DeploymentStats: time.Minute,
}, usage)
counts := map[string]int64{"vscode": 1, "cursor": 2, "zed": 3, "future_ide": 4, "jetbrains": 5, "reconnecting_pty": 6}
agentStat := dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: now.Add(-2 * time.Minute),
Usage: usage,
RxBytes: 1,
TxBytes: 1,
ConnectionCount: 1,
ConnectionMedianLatencyMS: 10,
SessionCounts: dbgen.SessionCounts(t, counts),
})
// Both ticker functions: template build times and deployment stats.
tickerTrap.MustWait(ctx).MustRelease(ctx)
tickerTrap.MustWait(ctx).MustRelease(ctx)
clock.Advance(time.Minute).MustWait(ctx)
stat, ok := cache.DeploymentStats()
require.True(t, ok, "cache should be populated after refresh")
// Legacy totals fold cursor into VS Code and zed into SSH.
require.Equal(t, codersdk.SessionCountDeploymentStats{
Apps: map[string]codersdk.SessionCountApp{
"vscode": {Count: 1, DisplayName: "VS Code", Icon: "/icon/code.svg", Family: codersdk.AppFamilyVSCode},
"cursor": {Count: 2, DisplayName: "Cursor", Icon: "/icon/cursor.svg", Family: codersdk.AppFamilyVSCode},
"zed": {Count: 3, DisplayName: "Zed", Icon: "/icon/zed.svg", Family: codersdk.AppFamilySSH},
"future_ide": {Count: 4, DisplayName: "future_ide", Family: codersdk.AppFamilyUnknown},
"jetbrains": {Count: 5, DisplayName: "JetBrains", Icon: "/icon/jetbrains.svg", Family: codersdk.AppFamilyJetBrains},
"reconnecting_pty": {Count: 6, DisplayName: "Web Terminal", Icon: "/icon/terminal.svg", Family: codersdk.AppFamilyReconnectingPTY},
},
VSCode: 3, SSH: 3, JetBrains: 5, ReconnectingPTY: 6,
}, stat.SessionCount)
// A later report with no sessions clears every app.
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
AgentID: agentStat.AgentID,
UserID: agentStat.UserID,
WorkspaceID: agentStat.WorkspaceID,
TemplateID: agentStat.TemplateID,
CreatedAt: now.Add(-time.Minute),
Usage: usage,
ConnectionMedianLatencyMS: 10,
SessionCounts: dbgen.SessionCounts(t, map[string]int64{}),
})
clock.Advance(time.Minute).MustWait(ctx)
empty, ok := cache.DeploymentStats()
require.True(t, ok)
require.Equal(t, codersdk.SessionCountDeploymentStats{
Apps: map[string]codersdk.SessionCountApp{},
}, empty.SessionCount)
require.Len(t, stat.SessionCount.Apps, len(counts), "refresh must not mutate a published snapshot")
})
}
}