Skip to content

Commit f4c1202

Browse files
mtojekmtojek
committed
Unit test
1 parent 5ebbfe9 commit f4c1202

File tree

3 files changed

+154
-24
lines changed

3 files changed

+154
-24
lines changed

‎cli/server.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,7 @@ func enablePrometheus(
199199
}
200200
afterCtx(ctx, closeWorkspacesFunc)
201201

202-
insightsMetricsCollector, err := insights.NewMetricsCollector(options.Database, options.Logger, 0)
202+
insightsMetricsCollector, err := insights.NewMetricsCollector(options.Database, options.Logger, 0, 0)
203203
if err != nil {
204204
return nil, xerrors.Errorf("unable to initialize insights metrics collector: %w", err)
205205
}

‎coderd/prometheusmetrics/insights/metricscollector.go

Lines changed: 30 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,10 @@ import (
1818
var templatesActiveUsersDesc = prometheus.NewDesc("coderd_insights_templates_active_users", "The number of active users of the template.", []string{"template_name"}, nil)
1919

2020
type MetricsCollector struct {
21-
database database.Store
22-
logger slog.Logger
23-
duration time.Duration
21+
database database.Store
22+
logger slog.Logger
23+
timeWindow time.Duration
24+
tickInterval time.Duration
2425

2526
data atomic.Pointer[insightsData]
2627
}
@@ -33,18 +34,22 @@ type insightsData struct {
3334

3435
var _ prometheus.Collector = new(MetricsCollector)
3536

36-
func NewMetricsCollector(db database.Store, logger slog.Logger, duration time.Duration) (*MetricsCollector, error) {
37-
if duration == 0 {
38-
duration = 5 * time.Minute
37+
func NewMetricsCollector(db database.Store, logger slog.Logger, timeWindow time.Duration, tickInterval time.Duration) (*MetricsCollector, error) {
38+
if timeWindow == 0 {
39+
timeWindow = 5 * time.Minute
3940
}
40-
if duration < 5*time.Minute {
41-
return nil, xerrors.Errorf("refresh interval must be at least 5 mins")
41+
if timeWindow < 5*time.Minute {
42+
return nil, xerrors.Errorf("time window must be at least 5 mins")
43+
}
44+
if tickInterval == 0 {
45+
tickInterval = timeWindow
4246
}
4347

4448
return &MetricsCollector{
45-
database: db,
46-
logger: logger.Named("insights_metrics_collector"),
47-
duration: duration,
49+
database: db,
50+
logger: logger.Named("insights_metrics_collector"),
51+
timeWindow: timeWindow,
52+
tickInterval: tickInterval,
4853
}, nil
4954
}
5055

@@ -56,10 +61,10 @@ func (mc *MetricsCollector) Run(ctx context.Context) (func(), error) {
5661
// correct duration after executing once.
5762
ticker := time.NewTicker(time.Nanosecond)
5863
doTick := func() {
59-
defer ticker.Reset(mc.duration)
64+
defer ticker.Reset(mc.tickInterval)
6065

6166
now := time.Now()
62-
startTime := now.Add(-mc.duration)
67+
startTime := now.Add(-mc.timeWindow)
6368
endTime := now
6469

6570
// TODO collect iteration time
@@ -89,20 +94,22 @@ func (mc *MetricsCollector) Run(ctx context.Context) (func(), error) {
8994

9095
// Phase 2: Collect template IDs, and fetch relevant details
9196
templateIDs := uniqueTemplateIDs(templateInsights)
92-
templates, err := mc.database.GetTemplatesWithFilter(ctx, database.GetTemplatesWithFilterParams{
93-
IDs: templateIDs,
94-
})
95-
if err != nil {
96-
mc.logger.Error(ctx, "unable to fetch template details from database", slog.Error(err))
97-
return
98-
}
9997

100-
templateNames := onlyTemplateNames(templates)
98+
templateNames := map[uuid.UUID]string{}
99+
if len(templateIDs) > 0 {
100+
templates, err := mc.database.GetTemplatesWithFilter(ctx, database.GetTemplatesWithFilterParams{
101+
IDs: templateIDs,
102+
})
103+
if err != nil {
104+
mc.logger.Error(ctx, "unable to fetch template details from database", slog.Error(err))
105+
return
106+
}
107+
templateNames = onlyTemplateNames(templates)
108+
}
101109

102110
// Refresh the collector state
103111
mc.data.Store(&insightsData{
104-
templates: templateInsights,
105-
112+
templates: templateInsights,
106113
templateNames: templateNames,
107114
})
108115
}

‎coderd/prometheusmetrics/insights/metricscollector_test.go

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
package insights_test
2+
3+
import (
4+
"context"
5+
"io"
6+
"testing"
7+
"time"
8+
9+
"github.com/google/uuid"
10+
"github.com/prometheus/client_golang/prometheus"
11+
"github.com/stretchr/testify/assert"
12+
"github.com/stretchr/testify/require"
13+
14+
"cdr.dev/slog/sloggers/slogtest"
15+
"github.com/coder/coder/v2/agent/agenttest"
16+
"github.com/coder/coder/v2/coderd/coderdtest"
17+
"github.com/coder/coder/v2/coderd/database/dbtestutil"
18+
"github.com/coder/coder/v2/coderd/prometheusmetrics/insights"
19+
"github.com/coder/coder/v2/codersdk"
20+
"github.com/coder/coder/v2/provisioner/echo"
21+
"github.com/coder/coder/v2/testutil"
22+
)
23+
24+
func TestCollect_TemplateInsights(t *testing.T) {
25+
t.Parallel()
26+
27+
logger := slogtest.Make(t, nil)
28+
db, ps := dbtestutil.NewDB(t)
29+
30+
options := &coderdtest.Options{
31+
IncludeProvisionerDaemon: true,
32+
AgentStatsRefreshInterval: time.Millisecond * 100,
33+
Database: db,
34+
Pubsub: ps,
35+
}
36+
client := coderdtest.New(t, options)
37+
38+
// Given
39+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
40+
defer cancel()
41+
42+
// Initialize metrics collector
43+
mc, err := insights.NewMetricsCollector(db, logger.Named("metrics_collector"), 0, time.Millisecond)
44+
require.NoError(t, err)
45+
46+
registry := prometheus.NewRegistry()
47+
registry.Register(mc)
48+
49+
closeFunc, err := mc.Run(ctx)
50+
require.NoError(t, err)
51+
t.Cleanup(closeFunc)
52+
53+
// Create two users, one that will appear in the report and another that
54+
// won't (due to not having/using a workspace).
55+
user := coderdtest.CreateFirstUser(t, client)
56+
_, _ = coderdtest.CreateAnotherUser(t, client, user.OrganizationID)
57+
authToken := uuid.NewString()
58+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
59+
Parse: echo.ParseComplete,
60+
ProvisionPlan: echo.PlanComplete,
61+
ProvisionApply: echo.ProvisionApplyWithAgent(authToken),
62+
})
63+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
64+
require.Empty(t, template.BuildTimeStats[codersdk.WorkspaceTransitionStart])
65+
66+
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
67+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
68+
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
69+
70+
// Start an agent so that we can generate stats.
71+
_ = agenttest.New(t, client.URL, authToken)
72+
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
73+
74+
// Connect to the agent to generate usage/latency stats.
75+
conn, err := client.DialWorkspaceAgent(ctx, resources[0].Agents[0].ID, &codersdk.DialWorkspaceAgentOptions{
76+
Logger: logger.Named("client"),
77+
})
78+
require.NoError(t, err)
79+
defer conn.Close()
80+
81+
sshConn, err := conn.SSHClient(ctx)
82+
require.NoError(t, err)
83+
defer sshConn.Close()
84+
85+
sess, err := sshConn.NewSession()
86+
require.NoError(t, err)
87+
defer sess.Close()
88+
89+
r, w := io.Pipe()
90+
defer r.Close()
91+
defer w.Close()
92+
sess.Stdin = r
93+
sess.Stdout = io.Discard
94+
err = sess.Start("cat")
95+
require.NoError(t, err)
96+
97+
collected := map[string]int{}
98+
assert.Eventuallyf(t, func() bool {
99+
// When
100+
metrics, err := registry.Gather()
101+
require.NoError(t, err)
102+
103+
// Then
104+
for _, metric := range metrics {
105+
switch metric.GetName() {
106+
case "coderd_insights_templates_active_users":
107+
for _, m := range metric.Metric {
108+
collected[metric.GetName()] = int(m.Gauge.GetValue())
109+
}
110+
default:
111+
require.FailNowf(t, "unexpected metric collected", "metric: %s", metric.GetName())
112+
}
113+
}
114+
115+
return len(collected) > 0
116+
}, testutil.WaitMedium, testutil.IntervalFast, "template insights are missing")
117+
118+
// We got our latency metrics, close the connection.
119+
_ = sess.Close()
120+
_ = sshConn.Close()
121+
122+
require.EqualValues(t, nil, collected)
123+
}

0 commit comments

Comments
 (0)