@@ -6,7 +6,7 @@ import { deleteWorkspaceFixture, insertWorkspaceFixture } from '@sim/db/testing/
66 * and finalizing each occurrence once. A local HTTP server stands in for the worker; each POST
77 * for a turn is one leg, which ends the turn, fails it, or stays open until the test acts.
88 */
9- import { envFlagsMock , resetEnvFlagsMock } from '@sim/testing/mocks/env-flags.mock'
9+ import { envFlagsMock , resetEnvFlagsMock , setEnvFlags } from '@sim/testing/mocks/env-flags.mock'
1010import { featureFlagsMock , featureFlagsMockFns } from '@sim/testing/mocks/feature-flags.mock'
1111import { afterAll , beforeAll , beforeEach , describe , expect , it , vi } from 'vitest'
1212
@@ -108,6 +108,11 @@ vi.mock('@/lib/mothership/jobs/reconcile', async (importOriginal) => {
108108 return { ...actual , reconcileJobRuns : vi . fn ( actual . reconcileJobRuns ) }
109109} )
110110vi . mock ( '@/lib/core/config/feature-flags' , ( ) => featureFlagsMock )
111+ /** Spied so a test can answer the hosted usage gate without a billing ledger. */
112+ vi . mock ( '@/lib/billing/core/attribution' , async ( importOriginal ) => {
113+ const actual = await importOriginal < typeof import ( '@/lib/billing/core/attribution' ) > ( )
114+ return { ...actual , checkAttributedUsageLimits : vi . fn ( actual . checkAttributedUsageLimits ) }
115+ } )
111116/** Spied so a test can tell an automatic turn-off, which emails the job's people, from none. */
112117vi . mock ( '@/lib/workflows/schedules/disable-notifications' , ( ) => ( {
113118 notifyScheduleAutoDisabled : vi . fn ( async ( ) => { } ) ,
@@ -139,7 +144,10 @@ import {
139144import { sleep } from '@sim/utils/helpers'
140145import { generateId } from '@sim/utils/id'
141146import { eq , inArray , sql } from 'drizzle-orm'
142- import { resolveBillingAttribution } from '@/lib/billing/core/attribution'
147+ import {
148+ checkAttributedUsageLimits ,
149+ resolveBillingAttribution ,
150+ } from '@/lib/billing/core/attribution'
143151import {
144152 acquireLock ,
145153 closeRedisConnection ,
@@ -256,11 +264,21 @@ async function runsOf(scheduleId: string) {
256264 status : jobExecutionLogs . status ,
257265 outcome : jobExecutionLogs . outcome ,
258266 summary : jobExecutionLogs . summary ,
267+ dispatchFailures : sql <
268+ number | null
269+ > `(${ jobExecutionLogs . executionData } ->>'dispatchFailures')::int` ,
259270 } )
260271 . from ( jobExecutionLogs )
261272 . where ( eq ( jobExecutionLogs . scheduleId , scheduleId ) )
262273}
263274
275+ /** The usage gate's answer when the payer's usage can't be read. */
276+ const usageOutage = {
277+ isExceeded : true ,
278+ message : 'Usage could not be checked.' ,
279+ reason : 'usage_unavailable' as const ,
280+ }
281+
264282/** Makes a schedule due again now, as its next occurrence arriving would. */
265283async function dueAgain ( scheduleId : string ) {
266284 const dueAt = new Date ( Date . now ( ) - 1_000 )
@@ -290,6 +308,8 @@ async function reconcileUntilSettled(executionId: string) {
290308}
291309
292310afterAll ( async ( ) => {
311+ /** A test that failed before ending its stream must not hold the server open. */
312+ worker . server . closeAllConnections ( )
293313 await new Promise < void > ( ( resolve ) => worker . server . close ( ( ) => resolve ( ) ) )
294314 for ( const [ key , value ] of Object . entries ( inheritedEnv ) ) {
295315 if ( value === undefined ) delete process . env [ key ]
@@ -634,6 +654,55 @@ describe.runIf(Boolean(redisUrl))('the scheduled job lane', () => {
634654 await reconcileUntilSettled ( occurrenceId )
635655 } )
636656
657+ it ( 'retries an occurrence whose usage could not be read, and runs it once usage reads again' , async ( ) => {
658+ setEnvFlags ( { isHosted : true } )
659+ const { scheduleId, occurrenceId } = await seedJob ( )
660+ vi . mocked ( checkAttributedUsageLimits ) . mockResolvedValueOnce ( usageOutage )
661+ await runJobLane ( new Date ( ) )
662+ expect ( await runsOf ( scheduleId ) ) . toEqual ( [
663+ expect . objectContaining ( { status : 'queued' , outcome : null , dispatchFailures : 1 } ) ,
664+ ] )
665+
666+ vi . mocked ( checkAttributedUsageLimits ) . mockResolvedValueOnce ( { isExceeded : false } )
667+ await reconcileJobRuns ( new Date ( Date . now ( ) + JOB_DISPATCH_GRACE_MS + 1_000 ) )
668+ await reconcileUntilSettled ( occurrenceId )
669+ expect ( legsFor ( occurrenceId ) ) . toHaveLength ( 1 )
670+ expect ( await runsOf ( scheduleId ) ) . toEqual ( [
671+ expect . objectContaining ( { status : 'complete' , outcome : 'not_met' , dispatchFailures : 1 } ) ,
672+ ] )
673+ } )
674+
675+ it ( 'ends an occurrence whose usage stays unreadable at its last attempt, and skips one whose billing is blocked' , async ( ) => {
676+ setEnvFlags ( { isHosted : true } )
677+ const outage = await seedJob ( )
678+ await claimDueJobOccurrences ( new Date ( ) , 10 )
679+ for ( let attempt = 1 ; attempt <= 5 ; attempt ++ ) {
680+ vi . mocked ( checkAttributedUsageLimits ) . mockResolvedValueOnce ( usageOutage )
681+ expect ( await dispatchJobOccurrence ( outage . occurrenceId ) ) . toBe (
682+ attempt < 5 ? 'retry' : 'failed'
683+ )
684+ }
685+ expect ( await runsOf ( outage . scheduleId ) ) . toEqual ( [
686+ expect . objectContaining ( { outcome : 'failed' , summary : 'The run could not be started.' } ) ,
687+ ] )
688+ expect ( legsFor ( outage . occurrenceId ) ) . toHaveLength ( 0 )
689+
690+ const blocked = await seedJob ( )
691+ await claimDueJobOccurrences ( new Date ( ) , 10 )
692+ vi . mocked ( checkAttributedUsageLimits ) . mockResolvedValueOnce ( {
693+ isExceeded : true ,
694+ message : 'Billing is blocked for this account.' ,
695+ reason : 'billing_blocked' ,
696+ } )
697+ expect ( await dispatchJobOccurrence ( blocked . occurrenceId ) ) . toBe ( 'skipped_budget' )
698+ expect ( await runsOf ( blocked . scheduleId ) ) . toEqual ( [
699+ expect . objectContaining ( {
700+ outcome : 'skipped_budget' ,
701+ summary : 'Billing is blocked for this account.' ,
702+ } ) ,
703+ ] )
704+ } )
705+
637706 it ( 'runs a one-shot job once and completes it' , async ( ) => {
638707 const { scheduleId, occurrenceId } = await seedJob ( { cronExpression : null } )
639708 await runJobLane ( new Date ( ) )
0 commit comments