@@ -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,6 +264,9 @@ 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 ) )
@@ -290,6 +301,8 @@ async function reconcileUntilSettled(executionId: string) {
290301}
291302
292303afterAll ( async ( ) => {
304+ /** A test that failed before ending its stream must not hold the server open. */
305+ worker . server . closeAllConnections ( )
293306 await new Promise < void > ( ( resolve ) => worker . server . close ( ( ) => resolve ( ) ) )
294307 for ( const [ key , value ] of Object . entries ( inheritedEnv ) ) {
295308 if ( value === undefined ) delete process . env [ key ]
@@ -634,6 +647,33 @@ describe.runIf(Boolean(redisUrl))('the scheduled job lane', () => {
634647 await reconcileUntilSettled ( occurrenceId )
635648 } )
636649
650+ it ( 'retries an occurrence whose usage could not be read, and skips one whose billing is blocked' , async ( ) => {
651+ setEnvFlags ( { isHosted : true } )
652+ const { scheduleId, occurrenceId } = await seedJob ( )
653+ vi . mocked ( checkAttributedUsageLimits ) . mockResolvedValueOnce ( {
654+ isExceeded : true ,
655+ message : 'Usage could not be checked.' ,
656+ reason : 'usage_unavailable' ,
657+ } )
658+ await runJobLane ( new Date ( ) )
659+ expect ( await runsOf ( scheduleId ) ) . toEqual ( [
660+ expect . objectContaining ( { status : 'queued' , outcome : null , dispatchFailures : 1 } ) ,
661+ ] )
662+
663+ vi . mocked ( checkAttributedUsageLimits ) . mockResolvedValueOnce ( {
664+ isExceeded : true ,
665+ message : 'Billing is blocked for this account.' ,
666+ reason : 'billing_blocked' ,
667+ } )
668+ expect ( await dispatchJobOccurrence ( occurrenceId ) ) . toBe ( 'skipped_budget' )
669+ expect ( await runsOf ( scheduleId ) ) . toEqual ( [
670+ expect . objectContaining ( {
671+ outcome : 'skipped_budget' ,
672+ summary : 'Billing is blocked for this account.' ,
673+ } ) ,
674+ ] )
675+ } )
676+
637677 it ( 'runs a one-shot job once and completes it' , async ( ) => {
638678 const { scheduleId, occurrenceId } = await seedJob ( { cronExpression : null } )
639679 await runJobLane ( new Date ( ) )
0 commit comments