mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
747 lines
24 KiB
Go
747 lines
24 KiB
Go
package usage_test
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"go.uber.org/goleak"
|
|
"go.uber.org/mock/gomock"
|
|
|
|
"cdr.dev/slog/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/database/dbtestutil"
|
|
"github.com/coder/coder/v2/coderd/database/dbtime"
|
|
"github.com/coder/coder/v2/coderd/rbac"
|
|
"github.com/coder/coder/v2/coderd/usage/usagetypes"
|
|
"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
|
|
"github.com/coder/coder/v2/enterprise/coderd/usage"
|
|
"github.com/coder/coder/v2/testutil"
|
|
"github.com/coder/quartz"
|
|
)
|
|
|
|
func TestMain(m *testing.M) {
|
|
goleak.VerifyTestMain(m, testutil.GoleakOptions...)
|
|
}
|
|
|
|
// TestIntegration tests the inserter and publisher by running them with a real
|
|
// database.
|
|
func TestIntegration(t *testing.T) {
|
|
t.Parallel()
|
|
const eventCount = 3
|
|
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
log := slogtest.Make(t, nil)
|
|
db, _ := dbtestutil.NewDB(t)
|
|
|
|
clock := quartz.NewMock(t)
|
|
deploymentID, licenseJWT := configureDeployment(ctx, t, db)
|
|
now := time.Now()
|
|
|
|
var (
|
|
calls int
|
|
handler func(req usagetypes.TallymanV1IngestRequest) any
|
|
)
|
|
ingestURL := fakeServer(t, tallymanHandler(t, deploymentID.String(), licenseJWT, func(req usagetypes.TallymanV1IngestRequest) any {
|
|
calls++
|
|
t.Logf("tallyman backend received call %d", calls)
|
|
|
|
if handler == nil {
|
|
t.Errorf("handler is nil")
|
|
return usagetypes.TallymanV1IngestResponse{}
|
|
}
|
|
return handler(req)
|
|
}))
|
|
|
|
inserter := usage.NewDBInserter(
|
|
usage.InserterWithClock(clock),
|
|
)
|
|
// Insert an old event that should never be published.
|
|
clock.Set(now.Add(-31 * 24 * time.Hour))
|
|
err := inserter.InsertDiscreteUsageEvent(ctx, db, usagetypes.DCManagedAgentsV1{
|
|
Count: 31,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Insert the events we expect to be published.
|
|
clock.Set(now.Add(1 * time.Second))
|
|
for i := 0; i < eventCount; i++ {
|
|
clock.Advance(time.Second)
|
|
err := inserter.InsertDiscreteUsageEvent(ctx, db, usagetypes.DCManagedAgentsV1{
|
|
Count: uint64(i + 1), // nolint:gosec // these numbers are tiny and will not overflow
|
|
})
|
|
require.NoErrorf(t, err, "collecting event %d", i)
|
|
}
|
|
|
|
// Wrap the publisher's DB in a dbauthz to ensure that the publisher has
|
|
// enough permissions.
|
|
authzDB := dbauthz.New(db, rbac.NewAuthorizer(prometheus.NewRegistry()), log, coderdtest.AccessControlStorePointer())
|
|
publisher := usage.NewTallymanPublisher(ctx, log, authzDB, coderdenttest.Keys,
|
|
usage.PublisherWithClock(clock),
|
|
usage.PublisherWithIngestURL(ingestURL),
|
|
)
|
|
defer publisher.Close()
|
|
|
|
// Start the publisher with a trap.
|
|
tickerTrap := clock.Trap().NewTicker()
|
|
defer tickerTrap.Close()
|
|
startErr := make(chan error)
|
|
go func() {
|
|
err := publisher.Start()
|
|
testutil.AssertSend(ctx, t, startErr, err)
|
|
}()
|
|
tickerCall := tickerTrap.MustWait(ctx)
|
|
tickerCall.MustRelease(ctx)
|
|
// The initial duration will always be some time between 5m and 17m.
|
|
require.GreaterOrEqual(t, tickerCall.Duration, 5*time.Minute)
|
|
require.LessOrEqual(t, tickerCall.Duration, 17*time.Minute)
|
|
require.NoError(t, testutil.RequireReceive(ctx, t, startErr))
|
|
|
|
// Set up a trap for the ticker.Reset call.
|
|
tickerResetTrap := clock.Trap().TickerReset()
|
|
defer tickerResetTrap.Close()
|
|
|
|
// Configure the handler for the first publish. This handler will accept the
|
|
// first event, temporarily reject the second, and permanently reject the
|
|
// third.
|
|
var temporarilyRejectedEventID string
|
|
handler = func(req usagetypes.TallymanV1IngestRequest) any {
|
|
// On the first call, accept the first event, temporarily reject the
|
|
// second, and permanently reject the third.
|
|
acceptedEvents := make([]usagetypes.TallymanV1IngestAcceptedEvent, 1)
|
|
rejectedEvents := make([]usagetypes.TallymanV1IngestRejectedEvent, 2)
|
|
if assert.Len(t, req.Events, eventCount) {
|
|
assert.JSONEqf(t, jsoninate(t, usagetypes.DCManagedAgentsV1{
|
|
Count: 1,
|
|
}), string(req.Events[0].EventData), "event data did not match for event %d", 0)
|
|
acceptedEvents[0].ID = req.Events[0].ID
|
|
|
|
temporarilyRejectedEventID = req.Events[1].ID
|
|
assert.JSONEqf(t, jsoninate(t, usagetypes.DCManagedAgentsV1{
|
|
Count: 2,
|
|
}), string(req.Events[1].EventData), "event data did not match for event %d", 1)
|
|
rejectedEvents[0].ID = req.Events[1].ID
|
|
rejectedEvents[0].Message = "temporarily rejected"
|
|
rejectedEvents[0].Permanent = false
|
|
|
|
assert.JSONEqf(t, jsoninate(t, usagetypes.DCManagedAgentsV1{
|
|
Count: 3,
|
|
}), string(req.Events[2].EventData), "event data did not match for event %d", 2)
|
|
rejectedEvents[1].ID = req.Events[2].ID
|
|
rejectedEvents[1].Message = "permanently rejected"
|
|
rejectedEvents[1].Permanent = true
|
|
}
|
|
return usagetypes.TallymanV1IngestResponse{
|
|
AcceptedEvents: acceptedEvents,
|
|
RejectedEvents: rejectedEvents,
|
|
}
|
|
}
|
|
|
|
// Advance the clock to the initial tick, which should trigger the first
|
|
// publish, then wait for the reset call. The duration will always be 17m
|
|
// for resets (only the initial tick is variable).
|
|
clock.Advance(tickerCall.Duration)
|
|
tickerResetCall := tickerResetTrap.MustWait(ctx)
|
|
require.Equal(t, 17*time.Minute, tickerResetCall.Duration)
|
|
tickerResetCall.MustRelease(ctx)
|
|
|
|
// The publisher should have published the events once.
|
|
require.Equal(t, 1, calls)
|
|
|
|
// Set the handler for the next publish call. This call should only include
|
|
// the temporarily rejected event from earlier. This time we'll accept it.
|
|
handler = func(req usagetypes.TallymanV1IngestRequest) any {
|
|
assert.Len(t, req.Events, 1)
|
|
acceptedEvents := make([]usagetypes.TallymanV1IngestAcceptedEvent, len(req.Events))
|
|
for i, event := range req.Events {
|
|
assert.Equal(t, temporarilyRejectedEventID, event.ID)
|
|
acceptedEvents[i].ID = event.ID
|
|
}
|
|
return usagetypes.TallymanV1IngestResponse{
|
|
AcceptedEvents: acceptedEvents,
|
|
RejectedEvents: []usagetypes.TallymanV1IngestRejectedEvent{},
|
|
}
|
|
}
|
|
|
|
// Advance the clock to the next tick and wait for the reset call.
|
|
clock.Advance(tickerResetCall.Duration)
|
|
tickerResetCall = tickerResetTrap.MustWait(ctx)
|
|
tickerResetCall.MustRelease(ctx)
|
|
|
|
// The publisher should have published the events again.
|
|
require.Equal(t, 2, calls)
|
|
|
|
// There should be no more publish calls after this, so set the handler to
|
|
// nil.
|
|
handler = nil
|
|
|
|
// Advance the clock to the next tick.
|
|
clock.Advance(tickerResetCall.Duration)
|
|
tickerResetTrap.MustWait(ctx).MustRelease(ctx)
|
|
|
|
// No publish should have taken place since there are no more events to
|
|
// publish.
|
|
require.Equal(t, 2, calls)
|
|
|
|
require.NoError(t, publisher.Close())
|
|
}
|
|
|
|
func TestPublisherNoEligibleLicenses(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
log := slogtest.Make(t, nil)
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
clock := quartz.NewMock(t)
|
|
|
|
// Configure the deployment manually.
|
|
deploymentID := uuid.New()
|
|
db.EXPECT().GetDeploymentID(gomock.Any()).Return(deploymentID.String(), nil).Times(1)
|
|
|
|
var calls int
|
|
ingestURL := fakeServer(t, tallymanHandler(t, deploymentID.String(), "", func(req usagetypes.TallymanV1IngestRequest) any {
|
|
calls++
|
|
return usagetypes.TallymanV1IngestResponse{
|
|
AcceptedEvents: []usagetypes.TallymanV1IngestAcceptedEvent{},
|
|
RejectedEvents: []usagetypes.TallymanV1IngestRejectedEvent{},
|
|
}
|
|
}))
|
|
|
|
publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys,
|
|
usage.PublisherWithClock(clock),
|
|
usage.PublisherWithIngestURL(ingestURL),
|
|
)
|
|
defer publisher.Close()
|
|
|
|
// Start the publisher with a trap.
|
|
tickerTrap := clock.Trap().NewTicker()
|
|
defer tickerTrap.Close()
|
|
startErr := make(chan error)
|
|
go func() {
|
|
err := publisher.Start()
|
|
testutil.RequireSend(ctx, t, startErr, err)
|
|
}()
|
|
tickerCall := tickerTrap.MustWait(ctx)
|
|
tickerCall.MustRelease(ctx)
|
|
require.NoError(t, testutil.RequireReceive(ctx, t, startErr))
|
|
|
|
// Mock zero licenses.
|
|
db.EXPECT().GetUnexpiredLicenses(gomock.Any()).Return([]database.License{}, nil).Times(1)
|
|
|
|
// Tick and wait for the reset call.
|
|
tickerResetTrap := clock.Trap().TickerReset()
|
|
defer tickerResetTrap.Close()
|
|
clock.Advance(tickerCall.Duration)
|
|
tickerResetCall := tickerResetTrap.MustWait(ctx)
|
|
tickerResetCall.MustRelease(ctx)
|
|
|
|
// The publisher should not have published the events.
|
|
require.Equal(t, 0, calls)
|
|
|
|
// Mock a single license with usage publishing disabled.
|
|
licenseJWT := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
|
PublishUsageData: false,
|
|
})
|
|
db.EXPECT().GetUnexpiredLicenses(gomock.Any()).Return([]database.License{
|
|
{
|
|
ID: 1,
|
|
JWT: licenseJWT,
|
|
UploadedAt: dbtime.Now(),
|
|
Exp: dbtime.Now().Add(48 * time.Hour), // fake
|
|
UUID: uuid.New(),
|
|
},
|
|
}, nil).Times(1)
|
|
|
|
// Tick and wait for the reset call.
|
|
clock.Advance(tickerResetCall.Duration)
|
|
tickerResetTrap.MustWait(ctx).MustRelease(ctx)
|
|
|
|
// The publisher should still not have published the events.
|
|
require.Equal(t, 0, calls)
|
|
}
|
|
|
|
// TestPublisherClaimExpiry tests the claim query to ensure that events are not
|
|
// claimed if they've recently been claimed by another publisher.
|
|
func TestPublisherClaimExpiry(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
log := slogtest.Make(t, nil)
|
|
db, _ := dbtestutil.NewDB(t)
|
|
clock := quartz.NewMock(t)
|
|
deploymentID, licenseJWT := configureDeployment(ctx, t, db)
|
|
now := time.Now()
|
|
|
|
var calls int
|
|
ingestURL := fakeServer(t, tallymanHandler(t, deploymentID.String(), licenseJWT, func(req usagetypes.TallymanV1IngestRequest) any {
|
|
calls++
|
|
return tallymanAcceptAllHandler(req)
|
|
}))
|
|
|
|
inserter := usage.NewDBInserter(
|
|
usage.InserterWithClock(clock),
|
|
)
|
|
|
|
publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys,
|
|
usage.PublisherWithClock(clock),
|
|
usage.PublisherWithIngestURL(ingestURL),
|
|
usage.PublisherWithInitialDelay(17*time.Minute),
|
|
)
|
|
defer publisher.Close()
|
|
|
|
// Create an event that was claimed 1h-18m ago. The ticker has a forced
|
|
// delay of 17m in this test.
|
|
clock.Set(now)
|
|
err := inserter.InsertDiscreteUsageEvent(ctx, db, usagetypes.DCManagedAgentsV1{
|
|
Count: 1,
|
|
})
|
|
require.NoError(t, err)
|
|
// Claim the event in the past. Claiming it this way via the database
|
|
// directly means it won't be marked as published or unclaimed.
|
|
events, err := db.SelectUsageEventsForPublishing(ctx, now.Add(-42*time.Minute))
|
|
require.NoError(t, err)
|
|
require.Len(t, events, 1)
|
|
|
|
// Start the publisher with a trap.
|
|
tickerTrap := clock.Trap().NewTicker()
|
|
defer tickerTrap.Close()
|
|
startErr := make(chan error)
|
|
go func() {
|
|
err := publisher.Start()
|
|
testutil.RequireSend(ctx, t, startErr, err)
|
|
}()
|
|
tickerCall := tickerTrap.MustWait(ctx)
|
|
require.Equal(t, 17*time.Minute, tickerCall.Duration)
|
|
tickerCall.MustRelease(ctx)
|
|
require.NoError(t, testutil.RequireReceive(ctx, t, startErr))
|
|
|
|
// Set up a trap for the ticker.Reset call.
|
|
tickerResetTrap := clock.Trap().TickerReset()
|
|
defer tickerResetTrap.Close()
|
|
|
|
// Advance the clock to the initial tick, which should trigger the first
|
|
// publish, then wait for the reset call. The duration will always be 17m
|
|
// for resets (only the initial tick is variable).
|
|
clock.Advance(tickerCall.Duration)
|
|
tickerResetCall := tickerResetTrap.MustWait(ctx)
|
|
require.Equal(t, 17*time.Minute, tickerResetCall.Duration)
|
|
tickerResetCall.MustRelease(ctx)
|
|
|
|
// No events should have been published since none are eligible.
|
|
require.Equal(t, 0, calls)
|
|
|
|
// Advance the clock to the next tick and wait for the reset call.
|
|
clock.Advance(tickerResetCall.Duration)
|
|
tickerResetCall = tickerResetTrap.MustWait(ctx)
|
|
tickerResetCall.MustRelease(ctx)
|
|
|
|
// The publisher should have published the event, as it's now eligible.
|
|
require.Equal(t, 1, calls)
|
|
}
|
|
|
|
// TestPublisherMissingEvents tests that the publisher notices events that are
|
|
// not returned by the Tallyman server and marks them as temporarily rejected.
|
|
func TestPublisherMissingEvents(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
log := slogtest.Make(t, nil)
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
deploymentID, licenseJWT := configureMockDeployment(t, db)
|
|
clock := quartz.NewMock(t)
|
|
now := time.Now()
|
|
clock.Set(now)
|
|
|
|
var calls int
|
|
ingestURL := fakeServer(t, tallymanHandler(t, deploymentID.String(), licenseJWT, func(req usagetypes.TallymanV1IngestRequest) any {
|
|
calls++
|
|
return usagetypes.TallymanV1IngestResponse{
|
|
AcceptedEvents: []usagetypes.TallymanV1IngestAcceptedEvent{},
|
|
RejectedEvents: []usagetypes.TallymanV1IngestRejectedEvent{},
|
|
}
|
|
}))
|
|
|
|
publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys,
|
|
usage.PublisherWithClock(clock),
|
|
usage.PublisherWithIngestURL(ingestURL),
|
|
)
|
|
|
|
// Expect the publisher to call SelectUsageEventsForPublishing, followed by
|
|
// UpdateUsageEventsPostPublish.
|
|
events := []database.UsageEvent{
|
|
{
|
|
ID: uuid.New().String(),
|
|
EventType: string(usagetypes.UsageEventTypeDCManagedAgentsV1),
|
|
EventData: []byte(jsoninate(t, usagetypes.DCManagedAgentsV1{
|
|
Count: 1,
|
|
})),
|
|
CreatedAt: now,
|
|
PublishedAt: sql.NullTime{},
|
|
PublishStartedAt: sql.NullTime{},
|
|
FailureMessage: sql.NullString{},
|
|
},
|
|
}
|
|
db.EXPECT().SelectUsageEventsForPublishing(gomock.Any(), gomock.Any()).Return(events, nil).Times(1)
|
|
db.EXPECT().UpdateUsageEventsPostPublish(gomock.Any(), gomock.Any()).DoAndReturn(
|
|
func(ctx context.Context, params database.UpdateUsageEventsPostPublishParams) error {
|
|
assert.Equal(t, []string{events[0].ID}, params.IDs)
|
|
assert.Equal(t, []string{"tallyman did not include the event in the response"}, params.FailureMessages)
|
|
assert.Equal(t, []bool{false}, params.SetPublishedAts)
|
|
return nil
|
|
},
|
|
).Times(1)
|
|
|
|
// Start the publisher with a trap.
|
|
tickerTrap := clock.Trap().NewTicker()
|
|
defer tickerTrap.Close()
|
|
startErr := make(chan error)
|
|
go func() {
|
|
err := publisher.Start()
|
|
testutil.RequireSend(ctx, t, startErr, err)
|
|
}()
|
|
tickerCall := tickerTrap.MustWait(ctx)
|
|
tickerCall.MustRelease(ctx)
|
|
require.NoError(t, testutil.RequireReceive(ctx, t, startErr))
|
|
|
|
// Tick and wait for the reset call.
|
|
tickerResetTrap := clock.Trap().TickerReset()
|
|
defer tickerResetTrap.Close()
|
|
clock.Advance(tickerCall.Duration)
|
|
tickerResetTrap.MustWait(ctx).MustRelease(ctx)
|
|
|
|
// The publisher should have published the events once.
|
|
require.Equal(t, 1, calls)
|
|
|
|
require.NoError(t, publisher.Close())
|
|
}
|
|
|
|
func TestPublisherLicenseSelection(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
log := slogtest.Make(t, nil)
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
clock := quartz.NewMock(t)
|
|
now := time.Now()
|
|
|
|
// Configure the deployment manually.
|
|
deploymentID := uuid.New()
|
|
db.EXPECT().GetDeploymentID(gomock.Any()).Return(deploymentID.String(), nil).Times(1)
|
|
|
|
// Insert multiple licenses:
|
|
// 1. PublishUsageData false, type=salesforce, iat 30m ago (ineligible, publish not enabled)
|
|
// 2. PublishUsageData true, type=trial, iat 1h ago (ineligible, not salesforce)
|
|
// 3. PublishUsageData true, type=salesforce, iat 30m ago, exp 10m ago (ineligible, expired)
|
|
// 4. PublishUsageData true, type=salesforce, iat 1h ago (eligible)
|
|
// 5. PublishUsageData true, type=salesforce, iat 30m ago (eligible, and newer!)
|
|
badLicense1 := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
|
PublishUsageData: false,
|
|
IssuedAt: now.Add(-30 * time.Minute),
|
|
})
|
|
badLicense2 := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
|
PublishUsageData: true,
|
|
IssuedAt: now.Add(-1 * time.Hour),
|
|
AccountType: "trial",
|
|
})
|
|
badLicense3 := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
|
PublishUsageData: true,
|
|
IssuedAt: now.Add(-30 * time.Minute),
|
|
ExpiresAt: now.Add(-10 * time.Minute),
|
|
})
|
|
badLicense4 := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
|
PublishUsageData: true,
|
|
IssuedAt: now.Add(-1 * time.Hour),
|
|
})
|
|
expectedLicense := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
|
PublishUsageData: true,
|
|
IssuedAt: now.Add(-30 * time.Minute),
|
|
})
|
|
// GetUnexpiredLicenses is not supposed to return expired licenses, but for
|
|
// the purposes of this test we're going to do it anyway.
|
|
db.EXPECT().GetUnexpiredLicenses(gomock.Any()).Return([]database.License{
|
|
{
|
|
ID: 1,
|
|
JWT: badLicense1,
|
|
Exp: now.Add(48 * time.Hour), // fake times, the JWT should be checked
|
|
UUID: uuid.New(),
|
|
UploadedAt: now,
|
|
},
|
|
{
|
|
ID: 2,
|
|
JWT: badLicense2,
|
|
Exp: now.Add(48 * time.Hour),
|
|
UUID: uuid.New(),
|
|
UploadedAt: now,
|
|
},
|
|
{
|
|
ID: 3,
|
|
JWT: badLicense3,
|
|
Exp: now.Add(48 * time.Hour),
|
|
UUID: uuid.New(),
|
|
UploadedAt: now,
|
|
},
|
|
{
|
|
ID: 4,
|
|
JWT: badLicense4,
|
|
Exp: now.Add(48 * time.Hour),
|
|
UUID: uuid.New(),
|
|
UploadedAt: now,
|
|
},
|
|
{
|
|
ID: 5,
|
|
JWT: expectedLicense,
|
|
Exp: now.Add(48 * time.Hour),
|
|
UUID: uuid.New(),
|
|
UploadedAt: now,
|
|
},
|
|
}, nil)
|
|
|
|
called := false
|
|
ingestURL := fakeServer(t, tallymanHandler(t, deploymentID.String(), expectedLicense, func(req usagetypes.TallymanV1IngestRequest) any {
|
|
called = true
|
|
return tallymanAcceptAllHandler(req)
|
|
}))
|
|
|
|
publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys,
|
|
usage.PublisherWithClock(clock),
|
|
usage.PublisherWithIngestURL(ingestURL),
|
|
)
|
|
defer publisher.Close()
|
|
|
|
// Start the publisher with a trap.
|
|
tickerTrap := clock.Trap().NewTicker()
|
|
defer tickerTrap.Close()
|
|
startErr := make(chan error)
|
|
go func() {
|
|
err := publisher.Start()
|
|
testutil.RequireSend(ctx, t, startErr, err)
|
|
}()
|
|
tickerCall := tickerTrap.MustWait(ctx)
|
|
tickerCall.MustRelease(ctx)
|
|
require.NoError(t, testutil.RequireReceive(ctx, t, startErr))
|
|
|
|
// Mock events to be published.
|
|
events := []database.UsageEvent{
|
|
{
|
|
ID: uuid.New().String(),
|
|
EventType: string(usagetypes.UsageEventTypeDCManagedAgentsV1),
|
|
EventData: []byte(jsoninate(t, usagetypes.DCManagedAgentsV1{
|
|
Count: 1,
|
|
})),
|
|
},
|
|
}
|
|
db.EXPECT().SelectUsageEventsForPublishing(gomock.Any(), gomock.Any()).Return(events, nil).Times(1)
|
|
db.EXPECT().UpdateUsageEventsPostPublish(gomock.Any(), gomock.Any()).DoAndReturn(
|
|
func(ctx context.Context, params database.UpdateUsageEventsPostPublishParams) error {
|
|
assert.Equal(t, []string{events[0].ID}, params.IDs)
|
|
assert.Equal(t, []string{""}, params.FailureMessages)
|
|
assert.Equal(t, []bool{true}, params.SetPublishedAts)
|
|
return nil
|
|
},
|
|
).Times(1)
|
|
|
|
// Tick and wait for the reset call.
|
|
tickerResetTrap := clock.Trap().TickerReset()
|
|
defer tickerResetTrap.Close()
|
|
clock.Advance(tickerCall.Duration)
|
|
tickerResetTrap.MustWait(ctx).MustRelease(ctx)
|
|
|
|
// The publisher should have published the events once.
|
|
require.True(t, called)
|
|
}
|
|
|
|
func TestPublisherTallymanError(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := testutil.Context(t, testutil.WaitLong)
|
|
log := slogtest.Make(t, nil)
|
|
ctrl := gomock.NewController(t)
|
|
db := dbmock.NewMockStore(ctrl)
|
|
clock := quartz.NewMock(t)
|
|
now := time.Now()
|
|
clock.Set(now)
|
|
|
|
deploymentID, licenseJWT := configureMockDeployment(t, db)
|
|
const errorMessage = "tallyman error"
|
|
var calls int
|
|
ingestURL := fakeServer(t, tallymanHandler(t, deploymentID.String(), licenseJWT, func(req usagetypes.TallymanV1IngestRequest) any {
|
|
calls++
|
|
return usagetypes.TallymanV1Response{
|
|
Message: errorMessage,
|
|
}
|
|
}))
|
|
|
|
publisher := usage.NewTallymanPublisher(ctx, log, db, coderdenttest.Keys,
|
|
usage.PublisherWithClock(clock),
|
|
usage.PublisherWithIngestURL(ingestURL),
|
|
)
|
|
defer publisher.Close()
|
|
|
|
// Start the publisher with a trap.
|
|
tickerTrap := clock.Trap().NewTicker()
|
|
defer tickerTrap.Close()
|
|
startErr := make(chan error)
|
|
go func() {
|
|
err := publisher.Start()
|
|
testutil.RequireSend(ctx, t, startErr, err)
|
|
}()
|
|
tickerCall := tickerTrap.MustWait(ctx)
|
|
tickerCall.MustRelease(ctx)
|
|
require.NoError(t, testutil.RequireReceive(ctx, t, startErr))
|
|
|
|
// Mock events to be published.
|
|
events := []database.UsageEvent{
|
|
{
|
|
ID: uuid.New().String(),
|
|
EventType: string(usagetypes.UsageEventTypeDCManagedAgentsV1),
|
|
EventData: []byte(jsoninate(t, usagetypes.DCManagedAgentsV1{
|
|
Count: 1,
|
|
})),
|
|
},
|
|
}
|
|
db.EXPECT().SelectUsageEventsForPublishing(gomock.Any(), gomock.Any()).Return(events, nil).Times(1)
|
|
db.EXPECT().UpdateUsageEventsPostPublish(gomock.Any(), gomock.Any()).DoAndReturn(
|
|
func(ctx context.Context, params database.UpdateUsageEventsPostPublishParams) error {
|
|
assert.Equal(t, []string{events[0].ID}, params.IDs)
|
|
assert.Contains(t, params.FailureMessages[0], errorMessage)
|
|
assert.Equal(t, []bool{false}, params.SetPublishedAts)
|
|
return nil
|
|
},
|
|
).Times(1)
|
|
|
|
// Tick and wait for the reset call.
|
|
tickerResetTrap := clock.Trap().TickerReset()
|
|
defer tickerResetTrap.Close()
|
|
clock.Advance(tickerCall.Duration)
|
|
tickerResetTrap.MustWait(ctx).MustRelease(ctx)
|
|
|
|
// The publisher should have published the events once.
|
|
require.Equal(t, 1, calls)
|
|
}
|
|
|
|
func jsoninate(t *testing.T, v any) string {
|
|
t.Helper()
|
|
if e, ok := v.(usagetypes.Event); ok {
|
|
v = e.Fields()
|
|
}
|
|
buf, err := json.Marshal(v)
|
|
require.NoError(t, err)
|
|
return string(buf)
|
|
}
|
|
|
|
func configureDeployment(ctx context.Context, t *testing.T, db database.Store) (uuid.UUID, string) {
|
|
t.Helper()
|
|
deploymentID := uuid.New()
|
|
err := db.InsertDeploymentID(ctx, deploymentID.String())
|
|
require.NoError(t, err)
|
|
|
|
licenseRaw := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
|
PublishUsageData: true,
|
|
})
|
|
_, err = db.InsertLicense(ctx, database.InsertLicenseParams{
|
|
UploadedAt: dbtime.Now(),
|
|
JWT: licenseRaw,
|
|
Exp: dbtime.Now().Add(48 * time.Hour),
|
|
UUID: uuid.New(),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
return deploymentID, licenseRaw
|
|
}
|
|
|
|
func configureMockDeployment(t *testing.T, db *dbmock.MockStore) (uuid.UUID, string) {
|
|
t.Helper()
|
|
deploymentID := uuid.New()
|
|
db.EXPECT().GetDeploymentID(gomock.Any()).Return(deploymentID.String(), nil).Times(1)
|
|
|
|
licenseRaw := coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
|
|
PublishUsageData: true,
|
|
})
|
|
db.EXPECT().GetUnexpiredLicenses(gomock.Any()).Return([]database.License{
|
|
{
|
|
ID: 1,
|
|
UploadedAt: dbtime.Now(),
|
|
JWT: licenseRaw,
|
|
Exp: dbtime.Now().Add(48 * time.Hour),
|
|
UUID: uuid.New(),
|
|
},
|
|
}, nil)
|
|
|
|
return deploymentID, licenseRaw
|
|
}
|
|
|
|
func fakeServer(t *testing.T, handler http.Handler) string {
|
|
t.Helper()
|
|
server := httptest.NewServer(handler)
|
|
t.Cleanup(server.Close)
|
|
return server.URL
|
|
}
|
|
|
|
func tallymanHandler(t *testing.T, expectDeploymentID string, expectLicenseJWT string, handler func(req usagetypes.TallymanV1IngestRequest) any) http.Handler {
|
|
t.Helper()
|
|
return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
|
|
t.Helper()
|
|
licenseJWT := r.Header.Get(usagetypes.TallymanCoderLicenseKeyHeader)
|
|
if expectLicenseJWT != "" && !assert.Equal(t, expectLicenseJWT, licenseJWT, "license JWT in request did not match") {
|
|
rw.WriteHeader(http.StatusUnauthorized)
|
|
_ = json.NewEncoder(rw).Encode(usagetypes.TallymanV1Response{
|
|
Message: "license JWT in request did not match",
|
|
})
|
|
return
|
|
}
|
|
|
|
deploymentID := r.Header.Get(usagetypes.TallymanCoderDeploymentIDHeader)
|
|
if expectDeploymentID != "" && !assert.Equal(t, expectDeploymentID, deploymentID, "deployment ID in request did not match") {
|
|
rw.WriteHeader(http.StatusUnauthorized)
|
|
_ = json.NewEncoder(rw).Encode(usagetypes.TallymanV1Response{
|
|
Message: "deployment ID in request did not match",
|
|
})
|
|
return
|
|
}
|
|
|
|
var req usagetypes.TallymanV1IngestRequest
|
|
err := json.NewDecoder(r.Body).Decode(&req)
|
|
if !assert.NoError(t, err, "could not decode request body") {
|
|
rw.WriteHeader(http.StatusBadRequest)
|
|
_ = json.NewEncoder(rw).Encode(usagetypes.TallymanV1Response{
|
|
Message: "could not decode request body",
|
|
})
|
|
return
|
|
}
|
|
|
|
resp := handler(req)
|
|
switch resp.(type) {
|
|
case usagetypes.TallymanV1Response:
|
|
rw.WriteHeader(http.StatusInternalServerError)
|
|
default:
|
|
rw.WriteHeader(http.StatusOK)
|
|
}
|
|
err = json.NewEncoder(rw).Encode(resp)
|
|
if !assert.NoError(t, err, "could not encode response body") {
|
|
rw.WriteHeader(http.StatusInternalServerError)
|
|
return
|
|
}
|
|
})
|
|
}
|
|
|
|
func tallymanAcceptAllHandler(req usagetypes.TallymanV1IngestRequest) usagetypes.TallymanV1IngestResponse {
|
|
acceptedEvents := make([]usagetypes.TallymanV1IngestAcceptedEvent, len(req.Events))
|
|
for i, event := range req.Events {
|
|
acceptedEvents[i].ID = event.ID
|
|
}
|
|
|
|
return usagetypes.TallymanV1IngestResponse{
|
|
AcceptedEvents: acceptedEvents,
|
|
RejectedEvents: []usagetypes.TallymanV1IngestRejectedEvent{},
|
|
}
|
|
}
|