mirror of
https://github.com/coder/coder.git
synced 2026-06-03 04:58:23 +00:00
Marcin Tojek
109 lines
3.4 KiB
Go
109 lines
3.4 KiB
Go
package usage_test
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"go.uber.org/mock/gomock"
|
|
|
|
"cdr.dev/slog/v3/sloggers/slogtest"
|
|
"github.com/coder/coder/v2/coderd/coderdtest"
|
|
"github.com/coder/coder/v2/coderd/database"
|
|
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
|
"github.com/coder/coder/v2/coderd/database/dbmock"
|
|
"github.com/coder/coder/v2/coderd/rbac"
|
|
"github.com/coder/coder/v2/coderd/usage/usagetypes"
|
|
"github.com/coder/coder/v2/enterprise/coderd/usage"
|
|
"github.com/coder/coder/v2/testutil"
|
|
"github.com/coder/quartz"
|
|
)
|
|
|
|
func TestCron(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
t.Run("BasicTick", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
clock := quartz.NewMock(t)
|
|
|
|
// The existence check should return false so the event gets
|
|
// inserted.
|
|
db.EXPECT().UsageEventExistsByID(gomock.Any(), gomock.Any()).
|
|
Return(false, nil).AnyTimes()
|
|
|
|
inserted := make(chan database.InsertUsageEventParams, 1)
|
|
db.EXPECT().InsertUsageEvent(gomock.Any(), gomock.Any()).
|
|
DoAndReturn(func(_ context.Context, params database.InsertUsageEventParams) error {
|
|
inserted <- params
|
|
return nil
|
|
}).AnyTimes()
|
|
|
|
inserter := usage.NewDBInserter(usage.InserterWithClock(clock))
|
|
cron := usage.NewCron(clock, slogtest.Make(t, nil), db, inserter)
|
|
require.NoError(t, cron.Register(usage.CronJob{
|
|
Name: "test-job",
|
|
Interval: 5 * time.Minute,
|
|
EventType: usagetypes.UsageEventTypeHBAISeatsV1,
|
|
Fn: func(_ context.Context) (usagetypes.HeartbeatEvent, error) {
|
|
return usagetypes.HBAISeats{Count: 42}, nil
|
|
},
|
|
}))
|
|
|
|
timerTrap := clock.Trap().NewTimer("test-job")
|
|
|
|
cron.Start(ctx)
|
|
defer cron.Close()
|
|
defer timerTrap.Close()
|
|
|
|
// Wait for timer creation, then fire it. The delay is the
|
|
// time until the next epoch-aligned boundary for the 5-minute
|
|
// interval — we don't assert the exact value since it depends
|
|
// on the mock clock's current time.
|
|
timerCall := timerTrap.MustWait(ctx)
|
|
timerCall.MustRelease(ctx)
|
|
clock.Advance(timerCall.Duration)
|
|
|
|
// Verify the event was inserted with an epoch-aligned ID.
|
|
select {
|
|
case params := <-inserted:
|
|
assert.Contains(t, params.ID, "hb_ai_seats_v1:")
|
|
case <-ctx.Done():
|
|
t.Fatal("timed out waiting for insert")
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestAISeatsHeartbeat checks that AISeatsHeartbeat returns the
|
|
// correct event type and count. It wraps a mock database with dbauthz
|
|
// to verify that the AsUsagePublisher subject has the required
|
|
// ResourceAiSeat.ActionRead permission.
|
|
func TestAISeatsHeartbeat(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
|
|
db.EXPECT().Wrappers().Return([]string{}).AnyTimes()
|
|
db.EXPECT().GetActiveAISeatCount(gomock.Any()).Return(int64(42), nil)
|
|
|
|
authz := rbac.NewStrictAuthorizer(prometheus.NewRegistry())
|
|
authzDB := dbauthz.New(db, authz, slogtest.Make(t, nil), coderdtest.AccessControlStorePointer())
|
|
|
|
// AISeatsHeartbeat internally uses AsUsagePublisher, which must
|
|
// have ResourceAiSeat.ActionRead to pass the dbauthz check.
|
|
fn := usage.AISeatsHeartbeat(authzDB)
|
|
event, err := fn(testutil.Context(t, testutil.WaitLong))
|
|
require.NoError(t, err)
|
|
|
|
hb, ok := event.(usagetypes.HBAISeats)
|
|
require.True(t, ok)
|
|
assert.Equal(t, int64(42), hb.Count)
|
|
}
|