fix: Fix Actor.charge behavior when the budget is overdrawn (#504)

janbuchar · web-flow · commit 7d17dfad5a90 · 2025-11-12T16:10:11.000+01:00
- Porting over the fix and tests from https://github.com/apify/apify-sdk-python/pulls/668
diff --git a/packages/apify/src/charging.ts b/packages/apify/src/charging.ts
@@ -425,15 +425,14 @@ export class ChargingManager {
             return Infinity;
         }
 
+        // The raw number of events allowed by the budget
+        const unroundedResult =
+            (this.maxTotalChargeUsd - this.calculateTotalChargedAmount()) /
+            price;
+
         // First round as Math.floor(4.9999999999999999) will incorrectly return 5
-        return Math.floor(
-            Number(
-                (
-                    (this.maxTotalChargeUsd -
-                        this.calculateTotalChargedAmount()) /
-                    price
-                ).toFixed(4),
-            ),
-        );
+        const roundedResult = Math.floor(Number(unroundedResult.toFixed(4)));
+
+        return Math.max(0, roundedResult);
     }
 }
diff --git a/test/MemoryStorageEmulator.ts b/test/MemoryStorageEmulator.ts
@@ -18,22 +18,30 @@ const LOCAL_EMULATION_DIR = resolve(
 
 export class MemoryStorageEmulator {
     protected localStorageDirectories: string[] = [];
+    protected storage?: MemoryStorage;
 
     async init(dirName = cryptoRandomObjectId(10)) {
         StorageManager.clearCache();
         const localStorageDir = resolve(LOCAL_EMULATION_DIR, dirName);
         this.localStorageDirectories.push(localStorageDir);
         await ensureDir(localStorageDir);
 
-        const storage = new MemoryStorage({
+        this.storage = new MemoryStorage({
             localDataDirectory: localStorageDir,
         });
-        Configuration.getGlobalConfig().useStorageClient(storage);
+        Configuration.getGlobalConfig().useStorageClient(this.storage);
         log.debug(
             `Initialized emulated memory storage in folder ${localStorageDir}`,
         );
     }
 
+    reapplyStorageClient() {
+        if (!this.storage) {
+            throw new Error('Storage not initialized. Call init() first.');
+        }
+        Configuration.getGlobalConfig().useStorageClient(this.storage);
+    }
+
     async destroy() {
         const promises = this.localStorageDirectories.map(async (dir) => {
             return rm(dir, { force: true, recursive: true });
diff --git a/test/apify/charging.test.ts b/test/apify/charging.test.ts
@@ -1,3 +1,5 @@
+import { afterEach, beforeEach, describe, expect, test } from 'vitest';
+
 import { Actor } from '../../packages/apify/src/index.js';
 import { MemoryStorageEmulator } from '../MemoryStorageEmulator.js';
 
@@ -24,6 +26,9 @@ describe('ChargingManager', () => {
         delete process.env.ACTOR_MAX_TOTAL_CHARGE_USD;
         delete process.env.APIFY_ACTOR_PRICING_INFO;
         delete process.env.APIFY_CHARGED_ACTOR_EVENT_COUNTS;
+        delete process.env.APIFY_IS_AT_HOME;
+        delete process.env.APIFY_TOKEN;
+        delete process.env.ACTOR_RUN_ID;
     });
 
     describe('charge()', () => {
@@ -66,6 +71,7 @@ describe('ChargingManager', () => {
             process.env.APIFY_CHARGED_ACTOR_EVENT_COUNTS = JSON.stringify({});
 
             await Actor.init();
+            localStorageEmulator.reapplyStorageClient();
 
             const chargeResult = await Actor.charge({
                 eventName: 'foobar',
@@ -75,5 +81,234 @@ describe('ChargingManager', () => {
             expect(chargeResult.chargedCount).toBe(4);
             expect(chargeResult.eventChargeLimitReached).toBe(false);
         });
+
+        test('should not call API when budget is exhausted during pushData', async () => {
+            // Don't use ACTOR_TEST_PAY_PER_EVENT when simulating Apify platform
+            delete process.env.ACTOR_TEST_PAY_PER_EVENT;
+
+            process.env.ACTOR_MAX_TOTAL_CHARGE_USD = '1.5';
+            process.env.APIFY_IS_AT_HOME = '1';
+            process.env.APIFY_TOKEN = 'this-wont-work';
+            process.env.ACTOR_RUN_ID = 'test-run-id';
+
+            // Set pricing info via env vars (as if coming from platform)
+            process.env.APIFY_ACTOR_PRICING_INFO = JSON.stringify({
+                pricingModel: 'PAY_PER_EVENT',
+                pricingPerEvent: {
+                    actorChargeEvents: {
+                        'some-event': {
+                            eventTitle: 'Some Event',
+                            eventPriceUsd: 1.0,
+                        },
+                        'another-event': {
+                            eventTitle: 'Another Event',
+                            eventPriceUsd: 1.0,
+                        },
+                    },
+                },
+            });
+            process.env.APIFY_CHARGED_ACTOR_EVENT_COUNTS = JSON.stringify({});
+
+            await Actor.init();
+            localStorageEmulator.reapplyStorageClient();
+
+            // Mock the API client charge method
+            const chargeSpy = vitest.fn().mockResolvedValue(undefined);
+            vitest.spyOn(Actor.apifyClient, 'run').mockReturnValue({
+                charge: chargeSpy,
+            } as any);
+
+            // Exhaust most of the budget (events cost $1 each)
+            const result1 = await Actor.charge({
+                eventName: 'some-event',
+                count: 1,
+            }); // Costs $1, leaving $0.5
+            expect(result1.chargedCount).toBe(1);
+            expect(chargeSpy).toHaveBeenCalledTimes(1);
+            expect(chargeSpy).toHaveBeenCalledWith({
+                eventName: 'some-event',
+                count: 1,
+            });
+
+            chargeSpy.mockClear();
+
+            // Now try to push data - we can't afford even 1 more event
+            const result = await Actor.pushData(
+                Array(10).fill({ hello: 'world' }),
+                'another-event',
+            );
+
+            // The API should NOT be called when count=0
+            expect(chargeSpy).not.toHaveBeenCalled();
+
+            // Correctly returns result with chargedCount=0
+            expect(result).toBeDefined();
+            expect(result!.chargedCount).toBe(0);
+            // Note: eventChargeLimitReached is false because count=0 was passed to charge()
+            // This is by design - see charging.ts line 318
+            expect(result!.eventChargeLimitReached).toBe(false);
+
+            // Verify no items were pushed (the important part)
+            const dataset = await Actor.openDataset();
+            const items = await dataset.getData();
+            expect(items.items).toHaveLength(0);
+
+            delete process.env.APIFY_IS_AT_HOME;
+            delete process.env.ACTOR_RUN_ID;
+        });
+
+        test('should verify charge API call with count=0 vs count>0', async () => {
+            // Don't use ACTOR_TEST_PAY_PER_EVENT when simulating Apify platform
+            delete process.env.ACTOR_TEST_PAY_PER_EVENT;
+
+            process.env.ACTOR_MAX_TOTAL_CHARGE_USD = '10.0';
+            process.env.APIFY_IS_AT_HOME = '1';
+            process.env.APIFY_TOKEN = 'this-wont-work';
+            process.env.ACTOR_RUN_ID = 'test-run-id';
+
+            // Set pricing info via env vars
+            process.env.APIFY_ACTOR_PRICING_INFO = JSON.stringify({
+                pricingModel: 'PAY_PER_EVENT',
+                pricingPerEvent: {
+                    actorChargeEvents: {
+                        'test-event': {
+                            eventTitle: 'Test Event',
+                            eventPriceUsd: 1.0,
+                        },
+                    },
+                },
+            });
+            process.env.APIFY_CHARGED_ACTOR_EVENT_COUNTS = JSON.stringify({});
+
+            await Actor.init();
+            localStorageEmulator.reapplyStorageClient();
+
+            // Mock the API client charge method
+            const chargeSpy = vitest.fn().mockResolvedValue(undefined);
+            vitest.spyOn(Actor.apifyClient, 'run').mockReturnValue({
+                charge: chargeSpy,
+            } as any);
+
+            // Call charge with count=0 - this should NOT call the API
+            const result1 = await Actor.charge({
+                eventName: 'test-event',
+                count: 0,
+            });
+            expect(chargeSpy).not.toHaveBeenCalled();
+            expect(result1.chargedCount).toBe(0);
+
+            // Call charge with count=1 - this SHOULD call the API
+            const result2 = await Actor.charge({
+                eventName: 'test-event',
+                count: 1,
+            });
+            expect(chargeSpy).toHaveBeenCalledTimes(1);
+            expect(chargeSpy).toHaveBeenCalledWith({
+                eventName: 'test-event',
+                count: 1,
+            });
+            expect(result2.chargedCount).toBe(1);
+        });
+    });
+
+    describe('calculateMaxEventChargeCountWithinLimit()', () => {
+        test('should not return negative values when budget is overdrawn', async () => {
+            // Don't use ACTOR_TEST_PAY_PER_EVENT when simulating Apify platform
+            delete process.env.ACTOR_TEST_PAY_PER_EVENT;
+
+            process.env.ACTOR_MAX_TOTAL_CHARGE_USD = '0.00025';
+
+            // Set pricing info and already charged events that overdraw the budget
+            process.env.APIFY_ACTOR_PRICING_INFO = JSON.stringify({
+                pricingModel: 'PAY_PER_EVENT',
+                pricingPerEvent: {
+                    actorChargeEvents: {
+                        event: {
+                            eventTitle: 'Event',
+                            eventPriceUsd: 0.0003,
+                        },
+                        'apify-actor-start': {
+                            eventTitle: 'Actor start',
+                            eventPriceUsd: 0.00005,
+                        },
+                    },
+                },
+            });
+            // Already charged 2 events worth $0.00035, which exceeds the $0.00025 budget
+            process.env.APIFY_CHARGED_ACTOR_EVENT_COUNTS = JSON.stringify({
+                event: 1,
+                'apify-actor-start': 1,
+            });
+
+            await Actor.init();
+
+            const chargingManager = Actor.getChargingManager();
+            const maxCount =
+                chargingManager.calculateMaxEventChargeCountWithinLimit(
+                    'event',
+                );
+
+            expect(maxCount).toBe(0);
+        });
+    });
+
+    describe('charge() with overdrawn budget', () => {
+        test('should handle charging when budget is already overdrawn', async () => {
+            // Don't use ACTOR_TEST_PAY_PER_EVENT when simulating Apify platform
+            delete process.env.ACTOR_TEST_PAY_PER_EVENT;
+
+            process.env.ACTOR_MAX_TOTAL_CHARGE_USD = '0.00025';
+            process.env.APIFY_IS_AT_HOME = '1';
+            process.env.APIFY_TOKEN = 'this-wont-work';
+            process.env.ACTOR_RUN_ID = 'test-run-id';
+
+            // Set pricing info
+            process.env.APIFY_ACTOR_PRICING_INFO = JSON.stringify({
+                pricingModel: 'PAY_PER_EVENT',
+                pricingPerEvent: {
+                    actorChargeEvents: {
+                        event: {
+                            eventTitle: 'Event',
+                            eventPriceUsd: 0.0003,
+                        },
+                        'apify-actor-start': {
+                            eventTitle: 'Actor start',
+                            eventPriceUsd: 0.00005,
+                        },
+                    },
+                },
+            });
+            // Already charged actor-start, leaving only $0.0002 which is less than the event cost
+            process.env.APIFY_CHARGED_ACTOR_EVENT_COUNTS = JSON.stringify({
+                event: 0,
+                'apify-actor-start': 1,
+            });
+
+            await Actor.init();
+            localStorageEmulator.reapplyStorageClient();
+
+            // Mock the API client charge method
+            const chargeSpy = vitest.fn().mockResolvedValue(undefined);
+            vitest.spyOn(Actor.apifyClient, 'run').mockReturnValue({
+                charge: chargeSpy,
+            } as any);
+
+            // Try to charge - the budget doesn't allow another event
+            const chargeResult = await Actor.charge({
+                eventName: 'event',
+                count: 1,
+            });
+            expect(chargeResult.chargedCount).toBe(0);
+
+            // Try to push data - the budget doesn't allow this either
+            const pushResult = await Actor.pushData(
+                [{ hello: 'world' }],
+                'event',
+            );
+            expect(pushResult!.chargedCount).toBe(0);
+
+            // The API should NOT have been called in either case
+            expect(chargeSpy).not.toHaveBeenCalled();
+        });
     });
 });
