opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: move usage types to new package (#19103)

Browse Source
This commit is contained in:
Dean Sheather
2025-08-20 23:48:38 +10:00
committed by GitHub
parent 6eb02d1c2a
commit 1a601c30ad
12 changed files with 470 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/coderd/usage/inserter.go
+8 -7
View File
@@ -10,20 +10,21 @@ import (
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbtime"
agplusage "github.com/coder/coder/v2/coderd/usage"
"github.com/coder/coder/v2/coderd/usage/usagetypes"
"github.com/coder/quartz"
)
// dbCollector collects usage events and stores them in the database for
// dbInserter collects usage events and stores them in the database for
// publishing.
type dbCollector struct {
type dbInserter struct {
clock quartz.Clock
}
var _ agplusage.Inserter = &dbCollector{}
var _ agplusage.Inserter = &dbInserter{}
// NewDBInserter creates a new database-backed usage event inserter.
func NewDBInserter(opts ...InserterOption) agplusage.Inserter {
c := &dbCollector{
c := &dbInserter{
clock: quartz.NewReal(),
}
for _, opt := range opts {
@@ -32,17 +33,17 @@ func NewDBInserter(opts ...InserterOption) agplusage.Inserter {
return c
}
type InserterOption func(*dbCollector)
type InserterOption func(*dbInserter)
// InserterWithClock sets the quartz clock to use for the inserter.
func InserterWithClock(clock quartz.Clock) InserterOption {
return func(c *dbCollector) {
return func(c *dbInserter) {
c.clock = clock
}
}
// InsertDiscreteUsageEvent implements agplusage.Inserter.
func (i *dbCollector) InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event agplusage.DiscreteEvent) error {
func (i *dbInserter) InsertDiscreteUsageEvent(ctx context.Context, tx database.Store, event usagetypes.DiscreteEvent) error {
if !event.EventType().IsDiscrete() {
return xerrors.Errorf("event type %q is not a discrete event", event.EventType())
}
enterprise/coderd/usage/inserter_test.go
+12 -12
View File
@@ -12,7 +12,7 @@ import (
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbmock"
"github.com/coder/coder/v2/coderd/database/dbtime"
agplusage "github.com/coder/coder/v2/coderd/usage"
"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"
@@ -33,37 +33,37 @@ func TestInserter(t *testing.T) {
now := dbtime.Now()
events := []struct {
time time.Time
event agplusage.DiscreteEvent
event usagetypes.DiscreteEvent
}{
{
time: now,
event: agplusage.DCManagedAgentsV1{
event: usagetypes.DCManagedAgentsV1{
Count: 1,
},
},
{
time: now.Add(1 * time.Minute),
event: agplusage.DCManagedAgentsV1{
event: usagetypes.DCManagedAgentsV1{
Count: 2,
},
},
}
for _, event := range events {
eventJSON := jsoninate(t, event.event)
for _, e := range events {
eventJSON := jsoninate(t, e.event)
db.EXPECT().InsertUsageEvent(gomock.Any(), gomock.Any()).DoAndReturn(
func(ctx any, params database.InsertUsageEventParams) error {
func(ctx interface{}, params database.InsertUsageEventParams) error {
_, err := uuid.Parse(params.ID)
assert.NoError(t, err)
assert.Equal(t, string(event.event.EventType()), params.EventType)
assert.Equal(t, e.event.EventType(), usagetypes.UsageEventType(params.EventType))
assert.JSONEq(t, eventJSON, string(params.EventData))
assert.Equal(t, event.time, params.CreatedAt)
assert.Equal(t, e.time, params.CreatedAt)
return nil
},
).Times(1)
clock.Set(event.time)
err := inserter.InsertDiscreteUsageEvent(ctx, db, event.event)
clock.Set(e.time)
err := inserter.InsertDiscreteUsageEvent(ctx, db, e.event)
require.NoError(t, err)
}
})
@@ -77,7 +77,7 @@ func TestInserter(t *testing.T) {
// We should get an error if the event is invalid.
inserter := usage.NewDBInserter()
err := inserter.InsertDiscreteUsageEvent(ctx, db, agplusage.DCManagedAgentsV1{
err := inserter.InsertDiscreteUsageEvent(ctx, db, usagetypes.DCManagedAgentsV1{
Count: 0, // invalid
})
assert.ErrorContains(t, err, `invalid "dc_managed_agents_v1" event: count must be greater than 0`)
enterprise/coderd/usage/publisher.go
+23 -56
View File
@@ -18,15 +18,13 @@ import (
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/pproflabel"
agplusage "github.com/coder/coder/v2/coderd/usage"
"github.com/coder/coder/v2/coderd/usage/usagetypes"
"github.com/coder/coder/v2/cryptorand"
"github.com/coder/coder/v2/enterprise/coderd/license"
"github.com/coder/quartz"
)
const (
CoderLicenseJWTHeader = "Coder-License-JWT"
tallymanURL = "https://tallyman-prod.coder.com"
tallymanIngestURLV1 = tallymanURL + "/api/v1/events/ingest"
@@ -217,20 +215,19 @@ func (p *tallymanPublisher) publishOnce(ctx context.Context, deploymentID uuid.U
var (
eventIDs = make(map[string]struct{})
tallymanReq = TallymanIngestRequestV1{
DeploymentID: deploymentID,
Events: make([]TallymanIngestEventV1, 0, len(events)),
tallymanReq = usagetypes.TallymanV1IngestRequest{
Events: make([]usagetypes.TallymanV1IngestEvent, 0, len(events)),
}
)
for _, event := range events {
eventIDs[event.ID] = struct{}{}
eventType := agplusage.EventType(event.EventType)
eventType := usagetypes.UsageEventType(event.EventType)
if !eventType.Valid() {
// This should never happen due to the check constraint in the
// database.
return 0, xerrors.Errorf("event %q has an invalid event type %q", event.ID, event.EventType)
}
tallymanReq.Events = append(tallymanReq.Events, TallymanIngestEventV1{
tallymanReq.Events = append(tallymanReq.Events, usagetypes.TallymanV1IngestEvent{
ID: event.ID,
EventType: eventType,
EventData: event.EventData,
@@ -243,17 +240,17 @@ func (p *tallymanPublisher) publishOnce(ctx context.Context, deploymentID uuid.U
return 0, xerrors.Errorf("duplicate event IDs found in events for publishing")
}
resp, err := p.sendPublishRequest(ctx, licenseJwt, tallymanReq)
resp, err := p.sendPublishRequest(ctx, deploymentID, licenseJwt, tallymanReq)
allFailed := err != nil
if err != nil {
p.log.Warn(ctx, "failed to send publish request to tallyman", slog.F("count", len(events)), slog.Error(err))
// Fake a response with all events temporarily rejected.
resp = TallymanIngestResponseV1{
AcceptedEvents: []TallymanIngestAcceptedEventV1{},
RejectedEvents: make([]TallymanIngestRejectedEventV1, len(events)),
resp = usagetypes.TallymanV1IngestResponse{
AcceptedEvents: []usagetypes.TallymanV1IngestAcceptedEvent{},
RejectedEvents: make([]usagetypes.TallymanV1IngestRejectedEvent, len(events)),
}
for i, event := range events {
resp.RejectedEvents[i] = TallymanIngestRejectedEventV1{
resp.RejectedEvents[i] = usagetypes.TallymanV1IngestRejectedEvent{
ID: event.ID,
Message: fmt.Sprintf("failed to publish to tallyman: %v", err),
Permanent: false,
@@ -267,8 +264,8 @@ func (p *tallymanPublisher) publishOnce(ctx context.Context, deploymentID uuid.U
p.log.Warn(ctx, "tallyman returned a different number of events than we sent", slog.F("sent", len(events)), slog.F("accepted", len(resp.AcceptedEvents)), slog.F("rejected", len(resp.RejectedEvents)))
}
acceptedEvents := make(map[string]*TallymanIngestAcceptedEventV1)
rejectedEvents := make(map[string]*TallymanIngestRejectedEventV1)
acceptedEvents := make(map[string]*usagetypes.TallymanV1IngestAcceptedEvent)
rejectedEvents := make(map[string]*usagetypes.TallymanV1IngestRejectedEvent)
for _, event := range resp.AcceptedEvents {
acceptedEvents[event.ID] = &event
}
@@ -389,37 +386,38 @@ func (p *tallymanPublisher) getBestLicenseJWT(ctx context.Context) (string, erro
return bestLicense.Raw, nil
}
func (p *tallymanPublisher) sendPublishRequest(ctx context.Context, licenseJwt string, req TallymanIngestRequestV1) (TallymanIngestResponseV1, error) {
func (p *tallymanPublisher) sendPublishRequest(ctx context.Context, deploymentID uuid.UUID, licenseJwt string, req usagetypes.TallymanV1IngestRequest) (usagetypes.TallymanV1IngestResponse, error) {
body, err := json.Marshal(req)
if err != nil {
return TallymanIngestResponseV1{}, err
return usagetypes.TallymanV1IngestResponse{}, err
}
r, err := http.NewRequestWithContext(ctx, http.MethodPost, p.ingestURL, bytes.NewReader(body))
if err != nil {
return TallymanIngestResponseV1{}, err
return usagetypes.TallymanV1IngestResponse{}, err
}
r.Header.Set(CoderLicenseJWTHeader, licenseJwt)
r.Header.Set(usagetypes.TallymanCoderLicenseKeyHeader, licenseJwt)
r.Header.Set(usagetypes.TallymanCoderDeploymentIDHeader, deploymentID.String())
resp, err := p.httpClient.Do(r)
if err != nil {
return TallymanIngestResponseV1{}, err
return usagetypes.TallymanV1IngestResponse{}, err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
var errBody TallymanErrorV1
var errBody usagetypes.TallymanV1Response
if err := json.NewDecoder(resp.Body).Decode(&errBody); err != nil {
errBody = TallymanErrorV1{
errBody = usagetypes.TallymanV1Response{
Message: fmt.Sprintf("could not decode error response body: %v", err),
}
}
return TallymanIngestResponseV1{}, xerrors.Errorf("unexpected status code %v, error: %s", resp.StatusCode, errBody.Message)
return usagetypes.TallymanV1IngestResponse{}, xerrors.Errorf("unexpected status code %v, error: %s", resp.StatusCode, errBody.Message)
}
var respBody TallymanIngestResponseV1
var respBody usagetypes.TallymanV1IngestResponse
if err := json.NewDecoder(resp.Body).Decode(&respBody); err != nil {
return TallymanIngestResponseV1{}, xerrors.Errorf("decode response body: %w", err)
return usagetypes.TallymanV1IngestResponse{}, xerrors.Errorf("decode response body: %w", err)
}
return respBody, nil
@@ -431,34 +429,3 @@ func (p *tallymanPublisher) Close() error {
<-p.done
return nil
}
type TallymanErrorV1 struct {
Message string `json:"message"`
}
type TallymanIngestRequestV1 struct {
DeploymentID uuid.UUID `json:"deployment_id"`
Events []TallymanIngestEventV1 `json:"events"`
}
type TallymanIngestEventV1 struct {
ID string `json:"id"`
EventType agplusage.EventType `json:"event_type"`
EventData json.RawMessage `json:"event_data"`
CreatedAt time.Time `json:"created_at"`
}
type TallymanIngestResponseV1 struct {
AcceptedEvents []TallymanIngestAcceptedEventV1 `json:"accepted_events"`
RejectedEvents []TallymanIngestRejectedEventV1 `json:"rejected_events"`
}
type TallymanIngestAcceptedEventV1 struct {
ID string `json:"id"`
}
type TallymanIngestRejectedEventV1 struct {
ID string `json:"id"`
Message string `json:"message"`
Permanent bool `json:"permanent"`
}
enterprise/coderd/usage/publisher_test.go
+69 -54
View File
@@ -24,7 +24,7 @@ import (
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/rbac"
agplusage "github.com/coder/coder/v2/coderd/usage"
"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"
@@ -51,16 +51,15 @@ func TestIntegration(t *testing.T) {
var (
calls int
handler func(req usage.TallymanIngestRequestV1) any
handler func(req usagetypes.TallymanV1IngestRequest) any
)
ingestURL := fakeServer(t, tallymanHandler(t, licenseJWT, func(req usage.TallymanIngestRequestV1) any {
ingestURL := fakeServer(t, tallymanHandler(t, deploymentID.String(), licenseJWT, func(req usagetypes.TallymanV1IngestRequest) any {
calls++
t.Logf("tallyman backend received call %d", calls)
assert.Equal(t, deploymentID, req.DeploymentID)
if handler == nil {
t.Errorf("handler is nil")
return usage.TallymanIngestResponseV1{}
return usagetypes.TallymanV1IngestResponse{}
}
return handler(req)
}))
@@ -70,7 +69,7 @@ func TestIntegration(t *testing.T) {
)
// Insert an old event that should never be published.
clock.Set(now.Add(-31 * 24 * time.Hour))
err := inserter.InsertDiscreteUsageEvent(ctx, db, agplusage.DCManagedAgentsV1{
err := inserter.InsertDiscreteUsageEvent(ctx, db, usagetypes.DCManagedAgentsV1{
Count: 31,
})
require.NoError(t, err)
@@ -79,7 +78,7 @@ func TestIntegration(t *testing.T) {
clock.Set(now.Add(1 * time.Second))
for i := 0; i < eventCount; i++ {
clock.Advance(time.Second)
err := inserter.InsertDiscreteUsageEvent(ctx, db, agplusage.DCManagedAgentsV1{
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)
@@ -117,33 +116,33 @@ func TestIntegration(t *testing.T) {
// first event, temporarily reject the second, and permanently reject the
// third.
var temporarilyRejectedEventID string
handler = func(req usage.TallymanIngestRequestV1) any {
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([]usage.TallymanIngestAcceptedEventV1, 1)
rejectedEvents := make([]usage.TallymanIngestRejectedEventV1, 2)
acceptedEvents := make([]usagetypes.TallymanV1IngestAcceptedEvent, 1)
rejectedEvents := make([]usagetypes.TallymanV1IngestRejectedEvent, 2)
if assert.Len(t, req.Events, eventCount) {
assert.JSONEqf(t, jsoninate(t, agplusage.DCManagedAgentsV1{
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, agplusage.DCManagedAgentsV1{
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, agplusage.DCManagedAgentsV1{
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 usage.TallymanIngestResponseV1{
return usagetypes.TallymanV1IngestResponse{
AcceptedEvents: acceptedEvents,
RejectedEvents: rejectedEvents,
}
@@ -162,16 +161,16 @@ func TestIntegration(t *testing.T) {
// 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 usage.TallymanIngestRequestV1) any {
handler = func(req usagetypes.TallymanV1IngestRequest) any {
assert.Len(t, req.Events, 1)
acceptedEvents := make([]usage.TallymanIngestAcceptedEventV1, len(req.Events))
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 usage.TallymanIngestResponseV1{
return usagetypes.TallymanV1IngestResponse{
AcceptedEvents: acceptedEvents,
RejectedEvents: []usage.TallymanIngestRejectedEventV1{},
RejectedEvents: []usagetypes.TallymanV1IngestRejectedEvent{},
}
}
@@ -211,11 +210,11 @@ func TestPublisherNoEligibleLicenses(t *testing.T) {
db.EXPECT().GetDeploymentID(gomock.Any()).Return(deploymentID.String(), nil).Times(1)
var calls int
ingestURL := fakeServer(t, tallymanHandler(t, "", func(req usage.TallymanIngestRequestV1) any {
ingestURL := fakeServer(t, tallymanHandler(t, deploymentID.String(), "", func(req usagetypes.TallymanV1IngestRequest) any {
calls++
return usage.TallymanIngestResponseV1{
AcceptedEvents: []usage.TallymanIngestAcceptedEventV1{},
RejectedEvents: []usage.TallymanIngestRejectedEventV1{},
return usagetypes.TallymanV1IngestResponse{
AcceptedEvents: []usagetypes.TallymanV1IngestAcceptedEvent{},
RejectedEvents: []usagetypes.TallymanV1IngestRejectedEvent{},
}
}))
@@ -280,11 +279,11 @@ func TestPublisherClaimExpiry(t *testing.T) {
log := slogtest.Make(t, nil)
db, _ := dbtestutil.NewDB(t)
clock := quartz.NewMock(t)
_, licenseJWT := configureDeployment(ctx, t, db)
deploymentID, licenseJWT := configureDeployment(ctx, t, db)
now := time.Now()
var calls int
ingestURL := fakeServer(t, tallymanHandler(t, licenseJWT, func(req usage.TallymanIngestRequestV1) any {
ingestURL := fakeServer(t, tallymanHandler(t, deploymentID.String(), licenseJWT, func(req usagetypes.TallymanV1IngestRequest) any {
calls++
return tallymanAcceptAllHandler(req)
}))
@@ -303,7 +302,7 @@ func TestPublisherClaimExpiry(t *testing.T) {
// 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, agplusage.DCManagedAgentsV1{
err := inserter.InsertDiscreteUsageEvent(ctx, db, usagetypes.DCManagedAgentsV1{
Count: 1,
})
require.NoError(t, err)
@@ -358,17 +357,17 @@ func TestPublisherMissingEvents(t *testing.T) {
log := slogtest.Make(t, nil)
ctrl := gomock.NewController(t)
db := dbmock.NewMockStore(ctrl)
_, licenseJWT := configureMockDeployment(t, db)
deploymentID, licenseJWT := configureMockDeployment(t, db)
clock := quartz.NewMock(t)
now := time.Now()
clock.Set(now)
var calls int
ingestURL := fakeServer(t, tallymanHandler(t, licenseJWT, func(req usage.TallymanIngestRequestV1) any {
ingestURL := fakeServer(t, tallymanHandler(t, deploymentID.String(), licenseJWT, func(req usagetypes.TallymanV1IngestRequest) any {
calls++
return usage.TallymanIngestResponseV1{
AcceptedEvents: []usage.TallymanIngestAcceptedEventV1{},
RejectedEvents: []usage.TallymanIngestRejectedEventV1{},
return usagetypes.TallymanV1IngestResponse{
AcceptedEvents: []usagetypes.TallymanV1IngestAcceptedEvent{},
RejectedEvents: []usagetypes.TallymanV1IngestRejectedEvent{},
}
}))
@@ -382,8 +381,8 @@ func TestPublisherMissingEvents(t *testing.T) {
events := []database.UsageEvent{
{
ID: uuid.New().String(),
EventType: string(agplusage.UsageEventTypeDCManagedAgentsV1),
EventData: []byte(jsoninate(t, agplusage.DCManagedAgentsV1{
EventType: string(usagetypes.UsageEventTypeDCManagedAgentsV1),
EventData: []byte(jsoninate(t, usagetypes.DCManagedAgentsV1{
Count: 1,
})),
CreatedAt: now,
@@ -508,9 +507,8 @@ func TestPublisherLicenseSelection(t *testing.T) {
}, nil)
called := false
ingestURL := fakeServer(t, tallymanHandler(t, expectedLicense, func(req usage.TallymanIngestRequestV1) any {
ingestURL := fakeServer(t, tallymanHandler(t, deploymentID.String(), expectedLicense, func(req usagetypes.TallymanV1IngestRequest) any {
called = true
assert.Equal(t, deploymentID, req.DeploymentID)
return tallymanAcceptAllHandler(req)
}))
@@ -536,8 +534,8 @@ func TestPublisherLicenseSelection(t *testing.T) {
events := []database.UsageEvent{
{
ID: uuid.New().String(),
EventType: string(agplusage.UsageEventTypeDCManagedAgentsV1),
EventData: []byte(jsoninate(t, agplusage.DCManagedAgentsV1{
EventType: string(usagetypes.UsageEventTypeDCManagedAgentsV1),
EventData: []byte(jsoninate(t, usagetypes.DCManagedAgentsV1{
Count: 1,
})),
},
@@ -572,12 +570,12 @@ func TestPublisherTallymanError(t *testing.T) {
now := time.Now()
clock.Set(now)
_, licenseJWT := configureMockDeployment(t, db)
deploymentID, licenseJWT := configureMockDeployment(t, db)
const errorMessage = "tallyman error"
var calls int
ingestURL := fakeServer(t, tallymanHandler(t, licenseJWT, func(req usage.TallymanIngestRequestV1) any {
ingestURL := fakeServer(t, tallymanHandler(t, deploymentID.String(), licenseJWT, func(req usagetypes.TallymanV1IngestRequest) any {
calls++
return usage.TallymanErrorV1{
return usagetypes.TallymanV1Response{
Message: errorMessage,
}
}))
@@ -604,8 +602,8 @@ func TestPublisherTallymanError(t *testing.T) {
events := []database.UsageEvent{
{
ID: uuid.New().String(),
EventType: string(agplusage.UsageEventTypeDCManagedAgentsV1),
EventData: []byte(jsoninate(t, agplusage.DCManagedAgentsV1{
EventType: string(usagetypes.UsageEventTypeDCManagedAgentsV1),
EventData: []byte(jsoninate(t, usagetypes.DCManagedAgentsV1{
Count: 1,
})),
},
@@ -632,7 +630,7 @@ func TestPublisherTallymanError(t *testing.T) {
func jsoninate(t *testing.T, v any) string {
t.Helper()
if e, ok := v.(agplusage.Event); ok {
if e, ok := v.(usagetypes.Event); ok {
v = e.Fields()
}
buf, err := json.Marshal(v)
@@ -688,44 +686,61 @@ func fakeServer(t *testing.T, handler http.Handler) string {
return server.URL
}
func tallymanHandler(t *testing.T, expectLicenseJWT string, handler func(req usage.TallymanIngestRequestV1) any) http.Handler {
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(usage.CoderLicenseJWTHeader)
licenseJWT := r.Header.Get(usagetypes.TallymanCoderLicenseKeyHeader)
if expectLicenseJWT != "" && !assert.Equal(t, expectLicenseJWT, licenseJWT, "license JWT in request did not match") {
rw.WriteHeader(http.StatusUnauthorized)
err := json.NewEncoder(rw).Encode(usage.TallymanErrorV1{
_ = json.NewEncoder(rw).Encode(usagetypes.TallymanV1Response{
Message: "license JWT in request did not match",
})
require.NoError(t, err)
return
}
var req usage.TallymanIngestRequestV1
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)
require.NoError(t, err)
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 usage.TallymanErrorV1:
case usagetypes.TallymanV1Response:
rw.WriteHeader(http.StatusInternalServerError)
default:
rw.WriteHeader(http.StatusOK)
}
err = json.NewEncoder(rw).Encode(resp)
require.NoError(t, err)
if !assert.NoError(t, err, "could not encode response body") {
rw.WriteHeader(http.StatusInternalServerError)
return
}
})
}
func tallymanAcceptAllHandler(req usage.TallymanIngestRequestV1) usage.TallymanIngestResponseV1 {
acceptedEvents := make([]usage.TallymanIngestAcceptedEventV1, len(req.Events))
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 usage.TallymanIngestResponseV1{
return usagetypes.TallymanV1IngestResponse{
AcceptedEvents: acceptedEvents,
RejectedEvents: []usage.TallymanIngestRejectedEventV1{},
RejectedEvents: []usagetypes.TallymanV1IngestRejectedEvent{},
}
}