Skip to content

Commit 28dd7d9

Browse files
GAlexIHUGAlexIHU
authored
fix: ent reset list logic (#4406)
1 parent e6577d5 commit 28dd7d9

2 files changed

Lines changed: 110 additions & 43 deletions

File tree

‎openmeter/entitlement/adapter/entitlement.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -771,11 +771,16 @@ func (a *entitlementDBAdapter) ListActiveEntitlementsWithExpiredUsagePeriod(ctx
771771
now := clock.Now()
772772

773773
query := withAllUsageResets(repo.db.Entitlement.Query(), params.Namespaces).
774-
Where(EntitlementActiveAt(params.Highwatermark)...).
774+
Where(EntitlementActiveAt(now)...).
775775
Where(
776+
db_entitlement.EntitlementTypeEQ(db_entitlement.EntitlementType(entitlement.EntitlementTypeMetered)),
776777
db_entitlement.CurrentUsagePeriodEndNotNil(),
777778
db_entitlement.CurrentUsagePeriodEndLTE(params.Highwatermark),
778779
db_entitlement.Or(db_entitlement.DeletedAtIsNil(), db_entitlement.DeletedAtGT(now)),
780+
db_entitlement.Or(
781+
db_entitlement.MeasureUsageFromNotNil(),
782+
db_entitlement.HasUsageReset(),
783+
),
779784
)
780785

781786
if len(params.Namespaces) > 0 {

‎openmeter/entitlement/adapter/entitlement_test.go‎

Lines changed: 104 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,12 @@ import (
3737
var m sync.Mutex
3838

3939
type deps struct {
40-
entRepo entitlement.EntitlementRepo
41-
featureRepo feature.FeatureRepo
42-
subjectRepo subject.Service
43-
customerRepo customer.Adapter
44-
meterRepo *meteradapter.Adapter
40+
entRepo entitlement.EntitlementRepo
41+
featureRepo feature.FeatureRepo
42+
subjectRepo subject.Service
43+
customerRepo customer.Adapter
44+
meterRepo *meteradapter.Adapter
45+
usageResetRepo meteredentitlement.UsageResetRepo
4546
}
4647

4748
func setup(t *testing.T) (deps deps, cleanup func()) {
@@ -74,6 +75,7 @@ func setup(t *testing.T) (deps deps, cleanup func()) {
7475
require.NotNilf(t, deps.meterRepo, "meter adapter must not be nil")
7576

7677
deps.entRepo = adapter.NewPostgresEntitlementRepo(dbClient)
78+
deps.usageResetRepo = adapter.NewPostgresUsageResetRepo(dbClient)
7779
deps.featureRepo = featureadapter.NewPostgresFeatureRepo(dbClient, logger)
7880

7981
// customer adapter for creating customers in tests
@@ -329,79 +331,139 @@ func TestListActiveEntitlementsWithExpiredUsagePeriod(t *testing.T) {
329331
ns := "ns1"
330332
featureKey := "feature1"
331333

332-
t.Run("Should return entitlements with expired usage period", func(t *testing.T) {
333-
ctx := context.Background()
334+
t.Run("Should return only resettable metered entitlements with expired usage period", func(t *testing.T) {
335+
ctx := t.Context()
334336
repo, cleanup := setup(t)
335337
defer cleanup()
336338

337-
// Let's create an example feature
338339
feature, err := repo.featureRepo.CreateFeature(ctx, feature.CreateFeatureInputs{
339340
Namespace: ns,
340341
Key: featureKey,
341342
Name: "Feature 1",
342343
})
343344
require.NoError(t, err)
344345

345-
// Let's set the current time
346346
now := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)
347347
clock.SetTime(now)
348348
defer clock.ResetTime()
349349

350-
// First, create the subjects
350+
expiredPeriod := &timeutil.ClosedPeriod{
351+
From: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
352+
To: time.Date(2025, 2, 1, 0, 0, 0, 0, time.UTC),
353+
}
354+
activePeriod := &timeutil.ClosedPeriod{
355+
From: time.Date(2025, 3, 1, 0, 0, 0, 0, time.UTC),
356+
To: time.Date(2025, 4, 1, 0, 0, 0, 0, time.UTC),
357+
}
358+
usagePeriod := lo.ToPtr(entitlement.NewUsagePeriodInputFromRecurrence(timeutil.Recurrence{
359+
Interval: timeutil.RecurrencePeriodMonth,
360+
Anchor: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
361+
}))
362+
351363
cust1 := createCustomerWithSubject(t, repo.subjectRepo, repo.customerRepo, ns, "subject1")
352364
cust2 := createCustomerWithSubject(t, repo.subjectRepo, repo.customerRepo, ns, "subject2")
365+
custStatic := createCustomerWithSubject(t, repo.subjectRepo, repo.customerRepo, ns, "subject-static")
366+
custBoolean := createCustomerWithSubject(t, repo.subjectRepo, repo.customerRepo, ns, "subject-boolean")
367+
custMalformed := createCustomerWithSubject(t, repo.subjectRepo, repo.customerRepo, ns, "subject-malformed")
368+
custResetAnchor := createCustomerWithSubject(t, repo.subjectRepo, repo.customerRepo, ns, "subject-reset-anchor")
353369

354-
// Then create two entitlements, one with expired usage period and one with no expired usage period
355370
ent1, err := repo.entRepo.CreateEntitlement(ctx, entitlement.CreateEntitlementRepoInputs{
356-
Namespace: ns,
357-
FeatureID: feature.ID,
358-
FeatureKey: featureKey,
359-
UsageAttribution: cust1.GetUsageAttribution(),
360-
EntitlementType: entitlement.EntitlementTypeMetered,
361-
MeasureUsageFrom: lo.ToPtr(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)),
362-
UsagePeriod: lo.ToPtr(entitlement.NewUsagePeriodInputFromRecurrence(timeutil.Recurrence{
363-
Interval: timeutil.RecurrencePeriodMonth,
364-
Anchor: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
365-
})),
366-
CurrentUsagePeriod: &timeutil.ClosedPeriod{
367-
From: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
368-
To: time.Date(2025, 2, 1, 0, 0, 0, 0, time.UTC),
369-
},
371+
Namespace: ns,
372+
FeatureID: feature.ID,
373+
FeatureKey: featureKey,
374+
UsageAttribution: cust1.GetUsageAttribution(),
375+
EntitlementType: entitlement.EntitlementTypeMetered,
376+
MeasureUsageFrom: lo.ToPtr(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)),
377+
UsagePeriod: usagePeriod,
378+
CurrentUsagePeriod: expiredPeriod,
370379
})
371380
require.NoError(t, err)
372381
require.NotNil(t, ent1)
373382

374383
ent2, err := repo.entRepo.CreateEntitlement(ctx, entitlement.CreateEntitlementRepoInputs{
375-
Namespace: ns,
376-
FeatureID: feature.ID,
377-
FeatureKey: featureKey,
378-
UsageAttribution: cust2.GetUsageAttribution(),
379-
EntitlementType: entitlement.EntitlementTypeMetered,
380-
MeasureUsageFrom: lo.ToPtr(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)),
381-
UsagePeriod: lo.ToPtr(entitlement.NewUsagePeriodInputFromRecurrence(timeutil.Recurrence{
382-
Interval: timeutil.RecurrencePeriodMonth,
383-
Anchor: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
384-
})),
385-
CurrentUsagePeriod: &timeutil.ClosedPeriod{
386-
From: time.Date(2025, 3, 1, 0, 0, 0, 0, time.UTC),
387-
To: time.Date(2025, 4, 1, 0, 0, 0, 0, time.UTC),
388-
},
384+
Namespace: ns,
385+
FeatureID: feature.ID,
386+
FeatureKey: featureKey,
387+
UsageAttribution: cust2.GetUsageAttribution(),
388+
EntitlementType: entitlement.EntitlementTypeMetered,
389+
MeasureUsageFrom: lo.ToPtr(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)),
390+
UsagePeriod: usagePeriod,
391+
CurrentUsagePeriod: activePeriod,
389392
})
390393
require.NoError(t, err)
391394
require.NotNil(t, ent2)
392395

396+
staticEnt, err := repo.entRepo.CreateEntitlement(ctx, entitlement.CreateEntitlementRepoInputs{
397+
Namespace: ns,
398+
FeatureID: feature.ID,
399+
FeatureKey: featureKey,
400+
UsageAttribution: custStatic.GetUsageAttribution(),
401+
EntitlementType: entitlement.EntitlementTypeStatic,
402+
Config: lo.ToPtr(`{"on":true}`),
403+
UsagePeriod: usagePeriod,
404+
CurrentUsagePeriod: expiredPeriod,
405+
})
406+
require.NoError(t, err)
407+
require.NotNil(t, staticEnt)
408+
409+
booleanEnt, err := repo.entRepo.CreateEntitlement(ctx, entitlement.CreateEntitlementRepoInputs{
410+
Namespace: ns,
411+
FeatureID: feature.ID,
412+
FeatureKey: featureKey,
413+
UsageAttribution: custBoolean.GetUsageAttribution(),
414+
EntitlementType: entitlement.EntitlementTypeBoolean,
415+
UsagePeriod: usagePeriod,
416+
CurrentUsagePeriod: expiredPeriod,
417+
})
418+
require.NoError(t, err)
419+
require.NotNil(t, booleanEnt)
420+
421+
malformedEnt, err := repo.entRepo.CreateEntitlement(ctx, entitlement.CreateEntitlementRepoInputs{
422+
Namespace: ns,
423+
FeatureID: feature.ID,
424+
FeatureKey: featureKey,
425+
UsageAttribution: custMalformed.GetUsageAttribution(),
426+
EntitlementType: entitlement.EntitlementTypeMetered,
427+
UsagePeriod: usagePeriod,
428+
CurrentUsagePeriod: expiredPeriod,
429+
})
430+
require.NoError(t, err)
431+
require.NotNil(t, malformedEnt)
432+
433+
resetAnchorEnt, err := repo.entRepo.CreateEntitlement(ctx, entitlement.CreateEntitlementRepoInputs{
434+
Namespace: ns,
435+
FeatureID: feature.ID,
436+
FeatureKey: featureKey,
437+
UsageAttribution: custResetAnchor.GetUsageAttribution(),
438+
EntitlementType: entitlement.EntitlementTypeMetered,
439+
UsagePeriod: usagePeriod,
440+
CurrentUsagePeriod: expiredPeriod,
441+
})
442+
require.NoError(t, err)
443+
require.NotNil(t, resetAnchorEnt)
444+
445+
err = repo.usageResetRepo.Save(ctx, meteredentitlement.UsageResetUpdate{
446+
NamespacedModel: models.NamespacedModel{
447+
Namespace: ns,
448+
},
449+
EntitlementID: resetAnchorEnt.ID,
450+
ResetTime: time.Date(2025, 1, 15, 0, 0, 0, 0, time.UTC),
451+
Anchor: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
452+
UsagePeriodInterval: timeutil.RecurrencePeriodMonth.ISOString(),
453+
})
454+
require.NoError(t, err)
455+
393456
now = time.Date(2025, 3, 1, 0, 0, 0, 0, time.UTC)
394457
clock.SetTime(now)
395-
defer clock.ResetTime()
396458

397-
// Let's check that the entitlement with expired usage period is returned
398459
ents, err := repo.entRepo.ListActiveEntitlementsWithExpiredUsagePeriod(ctx, entitlement.ListExpiredEntitlementsParams{
399460
Namespaces: []string{ns},
400461
Highwatermark: now,
401462
})
402463
require.NoError(t, err)
403-
require.Equal(t, 1, len(ents))
464+
require.Len(t, ents, 2)
404465
require.Equal(t, ent1.ID, ents[0].ID)
466+
require.Equal(t, resetAnchorEnt.ID, ents[1].ID)
405467
})
406468

407469
t.Run("Should return entitlements with cursor and limit", func(t *testing.T) {

0 commit comments

Comments
 (0)