Shade/packages/shade-storage-postgres/tests/postgres-prekey-store.test.ts

import { describe, test, expect, beforeAll, afterAll, beforeEach } from 'bun:test';
import postgres, { type Sql } from 'postgres';
import { PostgresPrekeyStore } from '../src/index.js';

function randBytes(n: number): Uint8Array {
  const buf = new Uint8Array(n);
  crypto.getRandomValues(buf);
  return buf;
}

const PG_URL = process.env.SHADE_TEST_PG_URL;

if (!PG_URL) {
  describe.skip('PostgresPrekeyStore (skipped — set SHADE_TEST_PG_URL)', () => {
    test('placeholder', () => {});
  });
} else {
  describe('PostgresPrekeyStore', () => {
    let sql: Sql;
    let store: PostgresPrekeyStore;

    beforeAll(async () => {
      sql = postgres(PG_URL!);
      await sql`DROP TABLE IF EXISTS shade_server_identities CASCADE`;
      await sql`DROP TABLE IF EXISTS shade_server_signed_prekeys CASCADE`;
      await sql`DROP TABLE IF EXISTS shade_server_one_time_prekeys CASCADE`;
      store = await PostgresPrekeyStore.fromClient(sql);
    });

    beforeEach(async () => {
      await sql`TRUNCATE shade_server_identities, shade_server_signed_prekeys, shade_server_one_time_prekeys`;
    });

    afterAll(async () => {
      await sql.end();
    });

    test('identity crud', async () => {
      const sig = randBytes(32);
      const dh = randBytes(32);
      await store.saveIdentity('bob', sig, dh);
      const id = await store.getIdentity('bob');
      expect(id!.identitySigningKey).toEqual(sig);
      expect(id!.identityDHKey).toEqual(dh);
    });

    test('signed prekey crud', async () => {
      const pub = randBytes(32);
      const sig = randBytes(64);
      await store.saveSignedPreKey('bob', 1, pub, sig);
      const spk = await store.getSignedPreKey('bob');
      expect(spk!.keyId).toBe(1);
      expect(spk!.publicKey).toEqual(pub);
      expect(spk!.signature).toEqual(sig);
    });

    test('one-time prekey FIFO consumption', async () => {
      await store.saveOneTimePreKeys('bob', [
        { keyId: 100, publicKey: randBytes(32) },
        { keyId: 101, publicKey: randBytes(32) },
        { keyId: 102, publicKey: randBytes(32) },
      ]);
      expect(await store.getOneTimePreKeyCount('bob')).toBe(3);

      const k1 = await store.consumeOneTimePreKey('bob');
      expect(k1!.keyId).toBe(100);

      const k2 = await store.consumeOneTimePreKey('bob');
      expect(k2!.keyId).toBe(101);

      const k3 = await store.consumeOneTimePreKey('bob');
      expect(k3!.keyId).toBe(102);

      expect(await store.consumeOneTimePreKey('bob')).toBeNull();
    });

    test('multi-address isolation', async () => {
      await store.saveIdentity('alice', randBytes(32), randBytes(32));
      await store.saveIdentity('bob', randBytes(32), randBytes(32));
      await store.saveOneTimePreKeys('alice', [{ keyId: 1, publicKey: randBytes(32) }]);
      await store.saveOneTimePreKeys('bob', [
        { keyId: 1, publicKey: randBytes(32) },
        { keyId: 2, publicKey: randBytes(32) },
      ]);

      expect(await store.getOneTimePreKeyCount('alice')).toBe(1);
      expect(await store.getOneTimePreKeyCount('bob')).toBe(2);

      await store.deleteAll('bob');

      expect(await store.getIdentity('alice')).not.toBeNull();
      expect(await store.getIdentity('bob')).toBeNull();
      expect(await store.getOneTimePreKeyCount('alice')).toBe(1);
      expect(await store.getOneTimePreKeyCount('bob')).toBe(0);
    });

    test('concurrent consume uses FOR UPDATE SKIP LOCKED', async () => {
      // Insert 10 keys
      const keys = Array.from({ length: 10 }, (_, i) => ({
        keyId: 1000 + i,
        publicKey: randBytes(32),
      }));
      await store.saveOneTimePreKeys('bob', keys);

      // Consume concurrently
      const consumed = await Promise.all(
        Array.from({ length: 10 }, () => store.consumeOneTimePreKey('bob')),
      );

      // All 10 should be consumed, each unique
      expect(consumed.filter((k) => k !== null).length).toBe(10);
      const ids = consumed.map((k) => k!.keyId).sort();
      const unique = new Set(ids);
      expect(unique.size).toBe(10); // no duplicates
    });
  });
}