netris-nestri/apps/api/test/billing-checkout-gpu.test.ts

import { createHmac } from 'node:crypto';

import { afterAll, afterEach, describe, expect, test } from 'bun:test';

import { AccessToken } from '@nestri/core/access-token/index';
import { Polar } from '@nestri/core/billing/polar';
import { Fixtures } from '@nestri/core/db/fixtures';
import { testDb } from '@nestri/core/db/test';
import { Env } from '@nestri/core/env';
import { Identifier } from '@nestri/core/id';
import { Machine } from '@nestri/core/machine/index';

import { app } from '../app/index';
import { TEST_FRONTEND_URL } from './setup';

const sql = testDb();
const createdUserIds: string[] = [];
const createdOrgIds: string[] = [];
const realFetch = globalThis.fetch;

const PAID = 'prod_paid';
const V710 = { model: 'AMD Radeon Pro V710', vramMib: 28672, pciId: '1002:7461' };
const RX = { model: 'AMD Radeon RX 9070', vramMib: 16384, pciId: '1002:7550' };
const H100 = { model: 'NVIDIA H100', vramMib: 81920, pciId: '10de:2330' };
const RTX = { model: 'NVIDIA RTX 4090', vramMib: 24576, pciId: '10de:2684' };

const secret = `whsec_${Buffer.from('d'.repeat(32)).toString('base64')}`;

function configure(organisationId?: string) {
	Env.init({
		NODE_ENV: 'test',
		FRONTEND_URL: TEST_FRONTEND_URL,
		POLAR_ACCESS_TOKEN: 'polar_test',
		POLAR_PRODUCT_ID: PAID,
		POLAR_WEBHOOK_SECRET: secret,
		...(organisationId ? { FLEET_ORGANISATION_ID: organisationId } : {})
	});
	Polar.reset();
}

async function person(label: string) {
	const owner = await Fixtures.owner(label);
	createdUserIds.push(owner.userId);
	const pat = await AccessToken.create({
		id: Identifier.ascending('accessToken'),
		ownerUserId: owner.userId,
		teamId: null,
		name: label
	});
	return {
		owner,
		headers: {
			authorization: `Bearer ${pat.token}`,
			'content-type': 'application/json'
		} as Record<string, string>
	};
}

async function fleet(gpus: (typeof V710)[][]) {
	const operator = await Fixtures.owner('operator');
	createdUserIds.push(operator.userId);
	const f = await Fixtures.fleet(operator, gpus.length);
	createdOrgIds.push(f.organisationId);
	for (const [i, g] of gpus.entries()) await Machine.setGpus({ id: f.machines[i]!.id, gpus: g });
	configure(f.organisationId);
	return f;
}

/** Stand in for Polar, remembering what was sent to create a checkout. */
function polarCheckout() {
	const sent: any[] = [];
	globalThis.fetch = (async (input: Request | string | URL, init?: RequestInit) => {
		const req = input instanceof Request ? input : new Request(String(input), init);
		if (req.url.includes('polar.sh')) {
			sent.push(await req.json());
			return new Response(JSON.stringify({ id: 'chk_1', url: 'https://polar.sh/pay/chk_1' }), {
				status: 201,
				headers: { 'content-type': 'application/json' }
			});
		}
		return realFetch(input as never, init);
	}) as typeof fetch;
	return sent;
}

let n = 0;
async function deliver(type: string, data: Record<string, unknown>) {
	const body = JSON.stringify({ type, data });
	const id = `msg_cg_${++n}`;
	const ts = Math.floor(Date.now() / 1000);
	// Standard Webhooks: HMAC-SHA256 over `id.timestamp.body`, keyed by the decoded secret.
	const signature = `v1,${createHmac('sha256', Buffer.from(secret.slice(6), 'base64'))
		.update(`${id}.${ts}.${body}`)
		.digest('base64')}`;
	return app.request('/billing/webhook', {
		method: 'POST',
		headers: {
			'webhook-id': id,
			'webhook-timestamp': String(ts),
			'webhook-signature': signature
		},
		body
	});
}

const paid = (teamId: string, sub: string, gpu?: string) => ({
	id: sub,
	productId: PAID,
	customer: { externalId: teamId },
	metadata: gpu ? { teamId, gpu } : { teamId }
});

const rooms = (userId: string) =>
	sql`select b.gpu_pci_id, b.label from box b where b.user_id = ${userId} and b.time_deleted is null`;

afterEach(() => {
	globalThis.fetch = realFetch;
	Polar.reset();
});

afterAll(async () => {
	await sql`delete from checkout_room where team_id in (select id from team where owner_id = any(${createdUserIds}))`;
	await sql`delete from box where user_id = any(${createdUserIds})`;
	await sql`delete from install_token where created_by_user_id = any(${createdUserIds})`;
	await sql`delete from machine where owner_user_id = any(${createdUserIds})`;
	await sql`delete from organisation where id = any(${createdOrgIds})`;
	await sql`delete from "user" where id = any(${createdUserIds})`;
});

describe('POST /billing/checkout with a GPU', () => {
	test('the model travels to the provider as checkout metadata', async () => {
		await fleet([[V710], [H100]]);
		const p = await person('co-meta');
		const sent = polarCheckout();
		const res = await app.request('/billing/checkout', {
			method: 'POST',
			headers: p.headers,
			body: JSON.stringify({ gpu: H100.model })
		});
		// The stand-in's reply is not a full checkout, so the SDK may refuse
		// to parse it; what matters here is what was sent.
		void res;
		expect(sent[0].metadata).toEqual({ teamId: p.owner.teamId, gpu: H100.model });
	});

	test('no gpu leaves the metadata as it was', async () => {
		await fleet([[V710]]);
		const p = await person('co-plain');
		const sent = polarCheckout();
		const res = await app.request('/billing/checkout', {
			method: 'POST',
			headers: p.headers,
			body: JSON.stringify({})
		});
		void res;
		expect(sent[0].metadata).toEqual({ teamId: p.owner.teamId });
	});

	test('a model the fleet does not have is refused, and nothing reaches the provider', async () => {
		await fleet([[V710]]);
		const p = await person('co-unknown');
		const sent = polarCheckout();
		const res = await app.request('/billing/checkout', {
			method: 'POST',
			headers: p.headers,
			body: JSON.stringify({ gpu: 'Imaginary 9000' })
		});
		expect(res.status).toBe(400);
		expect(sent).toHaveLength(0);
	});
});

describe('the webhook makes the room that was paid for', () => {
	test('a paid subscription with a gpu creates a room on that model', async () => {
		await fleet([[V710], [H100]]);
		const p = await person('wh-room');
		const res = await deliver('subscription.active', paid(p.owner.teamId, 'sub_wh_room', H100.model));
		expect(res.status).toBe(200);
		const made = await rooms(p.owner.userId);
		expect(made.map((r) => [r.gpu_pci_id, r.label])).toEqual([[H100.pciId, H100.model]]);
		const [row] = await sql`select substituted, placed_model from checkout_room where subscription_id = 'sub_wh_room'`;
		expect(row).toMatchObject({ substituted: false, placed_model: H100.model });
	});

	test('a redelivery, and the created event after the active one, add no second room', async () => {
		await fleet([[H100], [H100]]);
		const p = await person('wh-twice');
		const data = paid(p.owner.teamId, 'sub_wh_twice', H100.model);
		await deliver('subscription.created', data);
		await deliver('subscription.active', data);
		await deliver('subscription.active', data);
		expect(await rooms(p.owner.userId)).toHaveLength(1);
	});

	test('with the model taken, the closest free one of the same vendor is used and recorded', async () => {
		await fleet([[H100], [V710], [RX], [RTX]]);
		const taker = await person('wh-taker');
		await deliver('subscription.active', paid(taker.owner.teamId, 'sub_wh_taker', H100.model));
		const p = await person('wh-sub');
		// H100 is gone. Nvidia first, so the 4090 beats the closer-in-memory AMD cards.
		await deliver('subscription.active', paid(p.owner.teamId, 'sub_wh_sub', H100.model));
		const made = await rooms(p.owner.userId);
		expect(made.map((r) => r.gpu_pci_id)).toEqual([RTX.pciId]);
		const [row] = await sql`select substituted, requested_model, placed_model from checkout_room where subscription_id = 'sub_wh_sub'`;
		expect(row).toMatchObject({
			substituted: true,
			requested_model: H100.model,
			placed_model: RTX.model
		});
	});

	test('with nothing free, no room is made, the plan is still applied, and a retry stays quiet', async () => {
		await fleet([[V710]]);
		const taker = await person('wh-none-taker');
		await deliver('subscription.active', paid(taker.owner.teamId, 'sub_wh_nt', V710.model));
		const p = await person('wh-none');
		const data = paid(p.owner.teamId, 'sub_wh_none', V710.model);
		expect((await deliver('subscription.active', data)).status).toBe(200);
		expect(await rooms(p.owner.userId)).toHaveLength(0);
		const [team] = await sql`select plan from team where id = ${p.owner.teamId}`;
		expect(team!.plan).toBe('paid');
		await deliver('subscription.active', data);
		expect(await rooms(p.owner.userId)).toHaveLength(0);
	});

	test('a paid subscription with no gpu creates nothing', async () => {
		await fleet([[V710]]);
		const p = await person('wh-nogpu');
		await deliver('subscription.active', paid(p.owner.teamId, 'sub_wh_nogpu'));
		expect(await rooms(p.owner.userId)).toHaveLength(0);
	});
});