HiveGrid-Wireframes/hivegrid-cloud/apps/web/app/modules/admin-sso/jit-provisioning.server.test.ts

import { beforeEach, describe, expect, it, vi } from "vitest";
import type { AdminSsoClaims } from "./claims.server";
import { ADMIN_SSO_PROVIDER, jitProvisionAdmin, type JitProvisioningDbClient } from "./jit-provisioning.server";

type AdminRecord = {
  userId: string;
  ssoProvider: string | null;
  ssoSubjectId: string | null;
  ssoGroups: string[];
  lastSsoLoginAt: Date | null;
};

type UserRecord = { id: string; email: string; firstName: string; lastName: string };

function buildFakeDb() {
  const adminRows = new Map<string, AdminRecord>();
  const userRows = new Map<string, UserRecord>();
  let userCounter = 0;

  const db: JitProvisioningDbClient = {
    adminUser: {
      findUnique: vi.fn(async (args) => {
        const { ssoProvider, ssoSubjectId } = args.where.ssoProvider_ssoSubjectId;
        for (const row of adminRows.values()) {
          if (row.ssoProvider === ssoProvider && row.ssoSubjectId === ssoSubjectId) {
            return { userId: row.userId, ssoProvider: row.ssoProvider, ssoSubjectId: row.ssoSubjectId, ssoGroups: row.ssoGroups, lastSsoLoginAt: row.lastSsoLoginAt };
          }
        }
        return null;
      }),
      update: vi.fn(async (args) => {
        const existing = adminRows.get(args.where.userId);
        if (!existing) {
          const err = new Error("record not found") as Error & { code?: string };
          err.code = "P2025";
          throw err;
        }
        const updated: AdminRecord = {
          ...existing,
          ssoProvider: args.data.ssoProvider ?? existing.ssoProvider,
          ssoSubjectId: args.data.ssoSubjectId ?? existing.ssoSubjectId,
          ssoGroups: args.data.ssoGroups ?? existing.ssoGroups,
          lastSsoLoginAt: args.data.lastSsoLoginAt ?? existing.lastSsoLoginAt,
        };
        adminRows.set(args.where.userId, updated);
        return updated;
      }),
      create: vi.fn(async (args) => {
        const row: AdminRecord = {
          userId: args.data.userId,
          ssoProvider: args.data.ssoProvider,
          ssoSubjectId: args.data.ssoSubjectId,
          ssoGroups: [...args.data.ssoGroups],
          lastSsoLoginAt: args.data.lastSsoLoginAt,
        };
        adminRows.set(row.userId, row);
        return row;
      }),
    },
    user: {
      findUnique: vi.fn(async (args) => {
        for (const u of userRows.values()) {
          if (u.email === args.where.email) return { id: u.id, email: u.email };
        }
        return null;
      }),
      update: vi.fn(async (args) => {
        const existing = userRows.get(args.where.id);
        if (!existing) throw new Error(`user ${args.where.id} missing`);
        const updated: UserRecord = {
          ...existing,
          firstName: args.data.firstName ?? existing.firstName,
          lastName: args.data.lastName ?? existing.lastName,
        };
        userRows.set(args.where.id, updated);
        return updated;
      }),
      create: vi.fn(async (args) => {
        userCounter += 1;
        const id = `user_${userCounter}`;
        const row: UserRecord = {
          id,
          email: args.data.email,
          firstName: args.data.firstName,
          lastName: args.data.lastName,
        };
        userRows.set(id, row);
        return { id: row.id, email: row.email };
      }),
    },
  };

  return { db, adminRows, userRows };
}

function buildClaims(overrides: Partial<AdminSsoClaims> = {}): AdminSsoClaims {
  return {
    subject: "authentik-sub-1",
    email: "admin@example.test",
    emailVerified: true,
    name: "Admin Person",
    groups: ["toolctrl-admin"],
    ...overrides,
  };
}

describe("jitProvisionAdmin", () => {
  let env: ReturnType<typeof buildFakeDb>;

  beforeEach(() => {
    env = buildFakeDb();
  });

  it("creates User+AdminUser for a brand new toolctrl-admin claim and returns super_admin", async () => {
    const result = await jitProvisionAdmin({ claims: buildClaims() }, env.db);
    expect(result.ok).toBe(true);
    if (result.ok) {
      expect(result.admin.role).toBe("super_admin");
      expect(result.admin.created).toBe(true);
      expect(result.admin.email).toBe("admin@example.test");
    }
    expect(env.adminRows.size).toBe(1);
    expect(env.userRows.size).toBe(1);
    const admin = [...env.adminRows.values()][0]!;
    expect(admin.ssoProvider).toBe(ADMIN_SSO_PROVIDER);
    expect(admin.ssoSubjectId).toBe("authentik-sub-1");
    expect(admin.ssoGroups).toEqual(["toolctrl-admin"]);
  });

  it("maps toolctrl-staff to staff role", async () => {
    const result = await jitProvisionAdmin({ claims: buildClaims({ groups: ["toolctrl-staff"] }) }, env.db);
    expect(result.ok).toBe(true);
    if (result.ok) expect(result.admin.role).toBe("staff");
  });

  it("rejects an unknown-group claim with forbidden-group", async () => {
    const result = await jitProvisionAdmin({ claims: buildClaims({ groups: ["random"] }) }, env.db);
    expect(result.ok).toBe(false);
    if (!result.ok) expect(result.reason).toBe("forbidden-group");
    expect(env.adminRows.size).toBe(0);
    expect(env.userRows.size).toBe(0);
  });

  it("rejects an empty-group claim with forbidden-group", async () => {
    const result = await jitProvisionAdmin({ claims: buildClaims({ groups: [] }) }, env.db);
    expect(result.ok).toBe(false);
  });

  it("updates the existing SSO admin and refreshes groups + lastSsoLoginAt", async () => {
    // Seed: first login creates the admin.
    const t0 = new Date("2026-01-01T00:00:00Z");
    await jitProvisionAdmin({ claims: buildClaims(), now: t0 }, env.db);
    const admin = [...env.adminRows.values()][0]!;
    expect(admin.lastSsoLoginAt).toEqual(t0);

    // Second login with same sub but new groups + later timestamp.
    const t1 = new Date("2026-02-01T00:00:00Z");
    const result = await jitProvisionAdmin(
      { claims: buildClaims({ groups: ["toolctrl-staff"] }), now: t1 },
      env.db,
    );
    expect(result.ok).toBe(true);
    if (result.ok) {
      expect(result.admin.created).toBe(false);
      // Role re-syncs from current group membership (Authentik = source of truth).
      expect(result.admin.role).toBe("staff");
    }
    const refreshed = [...env.adminRows.values()][0]!;
    expect(refreshed.ssoGroups).toEqual(["toolctrl-staff"]);
    expect(refreshed.lastSsoLoginAt).toEqual(t1);
    // No new user row.
    expect(env.userRows.size).toBe(1);
  });

  it("links an existing email-matched legacy admin by attaching SSO fields", async () => {
    // Seed legacy admin (no SSO fields).
    env.userRows.set("legacy_1", { id: "legacy_1", email: "admin@example.test", firstName: "Old", lastName: "Admin" });
    env.adminRows.set("legacy_1", { userId: "legacy_1", ssoProvider: null, ssoSubjectId: null, ssoGroups: [], lastSsoLoginAt: null });

    const result = await jitProvisionAdmin({ claims: buildClaims() }, env.db);
    expect(result.ok).toBe(true);
    if (result.ok) {
      expect(result.admin.userId).toBe("legacy_1");
      expect(result.admin.linkedExisting).toBe(true);
      expect(result.admin.created).toBe(false);
    }
    const admin = env.adminRows.get("legacy_1")!;
    expect(admin.ssoProvider).toBe(ADMIN_SSO_PROVIDER);
    expect(admin.ssoSubjectId).toBe("authentik-sub-1");
    expect(env.userRows.size).toBe(1); // no duplicate user
  });

  it("creates an AdminUser when a legacy User exists but has no AdminUser row", async () => {
    env.userRows.set("legacy_2", { id: "legacy_2", email: "admin@example.test", firstName: "X", lastName: "Y" });
    // Note: no AdminUser row pre-existing.
    const result = await jitProvisionAdmin({ claims: buildClaims() }, env.db);
    expect(result.ok).toBe(true);
    expect(env.adminRows.size).toBe(1);
    expect(env.adminRows.get("legacy_2")?.ssoSubjectId).toBe("authentik-sub-1");
  });
});