agenthub/test/db-seed.test.ts

import { describe, it, expect, beforeAll, afterAll } from 'vitest';
import { Pool } from 'pg';
import { drizzle } from 'drizzle-orm/node-postgres';
import { agents, rooms, roomMembers } from '../src/db/schema.js';

describe('Database Schema and Seed', () => {
  let pool: Pool;
  let db: ReturnType<typeof drizzle>;

  beforeAll(() => {
    pool = new Pool({
      host: process.env.POSTGRES_HOST || 'localhost',
      port: Number(process.env.POSTGRES_PORT) || 5432,
      user: process.env.POSTGRES_USER || 'agenthub',
      password: process.env.POSTGRES_PASSWORD || 'agenthub',
      database: process.env.POSTGRES_DB || 'agenthub',
    });

    pool.on('connect', (client) => {
      client.query("SET TIME ZONE 'UTC'");
    });

    db = drizzle(pool);
  });

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

  it('should have seeded 3 agents', async () => {
    const result = await db.select().from(agents);
    expect(result).toHaveLength(3);

    const names = result.map((a) => a.name).sort();
    expect(names).toEqual(['alice', 'bob', 'cli']);

    const alice = result.find((a) => a.name === 'alice');
    expect(alice?.role).toBe('admin');
    expect(alice?.displayName).toBe('Alice (Admin)');

    const bob = result.find((a) => a.name === 'bob');
    expect(bob?.role).toBe('agent');

    const cli = result.find((a) => a.name === 'cli');
    expect(cli?.role).toBe('agent');
  });

  it('should have seeded 2 rooms', async () => {
    const result = await db.select().from(rooms);
    expect(result).toHaveLength(2);

    const slugs = result.map((r) => r.slug).sort();
    expect(slugs).toEqual(['general', 'incidents']);

    const general = result.find((r) => r.slug === 'general');
    expect(general?.name).toBe('General Discussion');

    const incidents = result.find((r) => r.slug === 'incidents');
    expect(incidents?.name).toBe('Incident Response');
  });

  it('should have seeded room memberships', async () => {
    const result = await db.select().from(roomMembers);
    expect(result.length).toBeGreaterThanOrEqual(5);

    const agentsList = await db.select().from(agents);
    const roomsList = await db.select().from(rooms);

    const alice = agentsList.find((a) => a.name === 'alice');
    const bob = agentsList.find((a) => a.name === 'bob');
    const cli = agentsList.find((a) => a.name === 'cli');
    const general = roomsList.find((r) => r.slug === 'general');
    const incidents = roomsList.find((r) => r.slug === 'incidents');

    expect(alice).toBeDefined();
    expect(bob).toBeDefined();
    expect(cli).toBeDefined();
    expect(general).toBeDefined();
    expect(incidents).toBeDefined();

    const aliceInGeneral = result.find((m) => m.agentId === alice?.id && m.roomId === general?.id);
    expect(aliceInGeneral).toBeDefined();

    const bobInGeneral = result.find((m) => m.agentId === bob?.id && m.roomId === general?.id);
    expect(bobInGeneral).toBeDefined();

    const cliInGeneral = result.find((m) => m.agentId === cli?.id && m.roomId === general?.id);
    expect(cliInGeneral).toBeDefined();
  });

  it('should have all tables created with correct structure', async () => {
    const tableQuery = await pool.query(
      `
      SELECT table_name
      FROM information_schema.tables
      WHERE table_schema = 'public'
        AND table_type = 'BASE TABLE'
      ORDER BY table_name
    `,
    );

    const tableNames = tableQuery.rows.map((r) => r.table_name);
    expect(tableNames).toContain('agents');
    expect(tableNames).toContain('api_tokens');
    expect(tableNames).toContain('rooms');
    expect(tableNames).toContain('room_members');
    expect(tableNames).toContain('messages');
    expect(tableNames).toContain('audit_events');
  });

  it('should enforce unique constraint on agent name', async () => {
    await expect(
      db.insert(agents).values({
        name: 'alice',
        displayName: 'Alice Duplicate',
        role: 'agent',
      }),
    ).rejects.toThrow();
  });

  it('should enforce unique constraint on room slug', async () => {
    const agentsList = await db.select().from(agents);
    const alice = agentsList.find((a) => a.name === 'alice');

    await expect(
      db.insert(rooms).values({
        slug: 'general',
        name: 'General Duplicate',
        createdBy: alice?.id,
      }),
    ).rejects.toThrow();
  });
});