puter/src/backend/drivers/integrationTestUtil.ts

/*
 * Copyright (C) 2024-present Puter Technologies Inc.
 *
 * This file is part of Puter.
 *
 * Puter is free software: you can redistribute it and/or modify it under the
 * terms of the GNU Affero General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option) any
 * later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
 * details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see
 * [https://www.gnu.org/licenses/](https://www.gnu.org/licenses/).
 */

/**
 * Shared utilities for AI provider integration tests.
 *
 * Each test reads its credentials from `PUTER_TEST_AI_*` env vars (loaded by
 * the vitest config's `PUTER_` prefix) and skips itself when the var is missing
 * — tests run only on developer machines and in CI environments that supply the
 * right secrets.
 *
 * Filename intentionally omits `.test.` so vitest does not treat this helper as
 * a test file.
 */

import type { EventClient } from '../clients/event/EventClient.js';
import { translateProviderError } from '../controllers/drivers/DriverController.js';
import type { Actor } from '../core/actor.js';
import { SYSTEM_ACTOR, makeActor } from '../core/actor.js';
import { runWithContext } from '../core/context.js';
import { NO_CREDIT_HOLD } from '../services/metering/types.js';
import type { DriverStreamResult } from './meta.js';
import type { AiMeteringService } from './util/aiCostFactor.js';

/**
 * Returns the env var value, or `undefined` if missing/empty. Used as the gate
 * for `describe.skipIf` blocks.
 */
export const optionalEnv = (name: string): string | undefined => {
    const v = process.env[name];
    return v && v.length > 0 ? v : undefined;
};

/**
 * Returns true when the env var is unset, signaling the test block should be
 * skipped. Pair with `describe.skipIf(skipUnlessEnv(...))`.
 */
export const skipUnlessEnv = (name: string): boolean => !optionalEnv(name);

/**
 * Per-test timeout for provider integration tests. The default 5s vitest
 * timeout is way too short for real API calls — image generation in particular
 * routinely takes 15–30s. Pass this as the third argument to `it(...)`.
 */
export const INTEGRATION_TEST_TIMEOUT_MS = 90_000;

/** Runs `fn` with every AI driver's cost factor hooked to `factor`. */
export const withAiCostFactorHook = async <T>(
    events: EventClient,
    factor: number,
    fn: () => Promise<T>,
): Promise<T> => {
    const hook = (_key: string, event: { factor: number }) => {
        event.factor = factor;
    };
    events.on('ai.cost.factor.*', hook);
    try {
        return await fn();
    } finally {
        events.off('ai.cost.factor.*', hook);
    }
};

/**
 * Returns a no-op MeteringService stub. Real metering would write to DynamoDB /
 * Redis, which integration tests for AI providers don't care about — we just
 * need the provider's metering calls to not throw and to short-circuit credit
 * checks.
 */
export const makeMeteringStub = (): AiMeteringService =>
    ({
        utilRecordUsageObject: () => Promise.resolve([] as never),
        incrementUsage: () => Promise.resolve({} as never),
        batchIncrementUsages: () => Promise.resolve([] as never),
        hasEnoughCredits: () => Promise.resolve(true),
        costFactor: () => Promise.resolve(1),
        reserveAiCredits: () => Promise.resolve(NO_CREDIT_HOLD),
        reserveCredits: () => Promise.resolve(NO_CREDIT_HOLD),
        getRemainingUsage: () => Promise.resolve(Number.MAX_SAFE_INTEGER),
        getReportedCosts: () => [],
    }) as unknown as AiMeteringService;

/** Reads a driver's stream result to its end, as the response pipe would. */
export const readStreamResult = async (result: unknown): Promise<Buffer> => {
    const chunks: Buffer[] = [];
    for await (const chunk of (result as DriverStreamResult).stream) {
        chunks.push(Buffer.from(chunk as Uint8Array));
    }
    return Buffer.concat(chunks);
};

/**
 * Run `fn` inside a request-scoped context with `SYSTEM_ACTOR` set, which is
 * what providers expect (`Context.get('actor')`). The system actor bypasses
 * metering / quota gates by design.
 */
export const withTestActor = <T>(
    fn: () => T | Promise<T>,
    actor: Actor = SYSTEM_ACTOR,
): Promise<T> =>
    Promise.resolve(
        runWithContext({ actor, requestId: 'integration-test' }, fn),
    );

/**
 * Runs `fn` as {@link withTestActor} and returns its failure the way
 * `/drivers/call` hands it to the caller.
 */
export const callerError = async (
    fn: () => Promise<unknown>,
    actor: Actor = SYSTEM_ACTOR,
): Promise<unknown> => {
    try {
        await withTestActor(fn, actor);
    } catch (e) {
        return translateProviderError(e);
    }
    throw new Error('expected the call to reject');
};

/**
 * The four actor shapes provider tests exercise: a direct user session, the
 * user's own app, an app-issued access token (attributed through
 * `effectiveApp`), and the system actor. Shared so identifiers are tested
 * identically across providers instead of copied per suite.
 */
export const makeActorMatrix = (): Actor[] => {
    const user = { id: 42, uuid: 'u42', username: 'alice' };
    const app = { uid: 'app-abc' };
    return [
        makeActor({ user }),
        makeActor({ user, app }),
        makeActor({
            user,
            accessToken: { uid: 'tok-1', issuer: makeActor({ user, app }) },
        }),
        SYSTEM_ACTOR,
    ];
};

/** Identifiers `makeActorMatrix()` should produce, in matrix order. */
export const ACTOR_MATRIX_IDENTIFIERS: (string | undefined)[] = [
    'puter-u42',
    'puter-u42-app-abc',
    'puter-u42-app-abc',
    undefined,
];

/**
 * Picks `fields` off the first argument of each recorded mock call, keyed by
 * field, so a suite can compare what a provider sent against
 * `expectedIdentifierFields(fields)` with one `toEqual`.
 */
export const sentIdentifierFields = (
    calls: unknown[][],
    fields: string[],
): Record<string, unknown[]> =>
    Object.fromEntries(
        fields.map((field) => [
            field,
            calls.map((call) => (call[0] as Record<string, unknown>)[field]),
        ]),
    );

export const expectedIdentifierFields = (
    fields: string[],
): Record<string, unknown[]> =>
    Object.fromEntries(
        fields.map((field) => [field, ACTOR_MATRIX_IDENTIFIERS]),
    );