xtablo-source/apps/main/src/contexts/UpgradeBlockContext.test.tsx

import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { render, screen } from "@testing-library/react";
import { describe, expect, it, vi } from "vitest";
import { TestUserStoreProvider, type User } from "../providers/UserStoreProvider";
import { UpgradeBlockProvider, useMaybeUpgradeBlock, useUpgradeBlock } from "./UpgradeBlockContext";

// Mock the organization hook
vi.mock("../hooks/organization", () => ({
  useOrganization: vi.fn(),
}));

import { useOrganization } from "../hooks/organization";

const mockUseOrganization = vi.mocked(useOrganization);

const baseUser: User = {
  id: "user-1",
  short_user_id: "u1",
  name: "Test User",
  first_name: "Test",
  last_name: "User",
  email: "test@example.com",
  avatar_url: null,
  is_temporary: false,
  is_client: false,
  client_onboarded_at: null,
  last_signed_in: null,
  plan: "none",
  created_at: new Date().toISOString(),
};

const queryClient = new QueryClient({
  defaultOptions: { queries: { retry: false } },
});

function TestConsumer() {
  const context = useUpgradeBlock();
  return (
    <div>
      <span data-testid="blocked">{String(context.isBlocked)}</span>
      <span data-testid="reason">{context.reason ?? "none"}</span>
    </div>
  );
}

function MaybeTestConsumer() {
  const context = useMaybeUpgradeBlock();
  return (
    <div>
      <span data-testid="has-context">{String(context !== null)}</span>
    </div>
  );
}

function renderWithUser(user: User | null, orgData: ReturnType<typeof useOrganization>["data"]) {
  mockUseOrganization.mockReturnValue({
    data: orgData,
    isLoading: false,
    error: null,
  } as unknown as ReturnType<typeof useOrganization>);

  return render(
    <QueryClientProvider client={queryClient}>
      <TestUserStoreProvider user={user}>
        <UpgradeBlockProvider>
          <TestConsumer />
        </UpgradeBlockProvider>
      </TestUserStoreProvider>
    </QueryClientProvider>
  );
}

const compliantOrgData = {
  organization: {
    id: 1,
    name: "Test Org",
    plan: "team",
    member_count: 2,
    tablo_count: 5,
    logo_url: null,
  },
  members: [],
  invites_sent: [],
  trial_starts_at: "2026-01-01",
  trial_ends_at: "2026-02-01",
  is_trial_expired: false,
  required_plan: "solo" as const,
  required_team_quantity: 1,
  active_subscription_plan: "team" as const,
  active_subscription_quantity: 2,
  is_billing_owner: true,
};

const noPlanOrgData = {
  ...compliantOrgData,
  active_subscription_plan: null,
  active_subscription_quantity: 0,
};

const trialExpiredOrgData = {
  ...compliantOrgData,
  is_trial_expired: true,
  required_plan: "team" as const,
  required_team_quantity: 5,
  active_subscription_plan: "team" as const,
  active_subscription_quantity: 2,
};

describe("UpgradeBlockProvider", () => {
  it("is not blocked when user and org are loading (null)", () => {
    renderWithUser(null, undefined);
    expect(screen.getByTestId("blocked").textContent).toBe("false");
    expect(screen.getByTestId("reason").textContent).toBe("none");
  });

  it("is not blocked when user is loaded but org data is still loading", () => {
    renderWithUser(baseUser, undefined);
    expect(screen.getByTestId("blocked").textContent).toBe("false");
  });

  it("is not blocked for temporary users regardless of org billing state", () => {
    const temporaryUser = { ...baseUser, is_temporary: true };
    renderWithUser(temporaryUser, noPlanOrgData);
    expect(screen.getByTestId("blocked").textContent).toBe("false");
    expect(screen.getByTestId("reason").textContent).toBe("none");
  });

  it("is not blocked for temporary users even with expired trial", () => {
    const temporaryUser = { ...baseUser, is_temporary: true };
    renderWithUser(temporaryUser, trialExpiredOrgData);
    expect(screen.getByTestId("blocked").textContent).toBe("false");
  });

  it("blocks regular users when org has no active plan", () => {
    renderWithUser(baseUser, noPlanOrgData);
    expect(screen.getByTestId("blocked").textContent).toBe("true");
    expect(screen.getByTestId("reason").textContent).toBe("no_plan");
  });

  it("blocks regular users when trial expired and plan is insufficient", () => {
    renderWithUser(baseUser, trialExpiredOrgData);
    expect(screen.getByTestId("blocked").textContent).toBe("true");
    expect(screen.getByTestId("reason").textContent).toBe("trial_expired");
  });

  it("is not blocked for regular users with a compliant subscription", () => {
    renderWithUser(baseUser, compliantOrgData);
    expect(screen.getByTestId("blocked").textContent).toBe("false");
    expect(screen.getByTestId("reason").textContent).toBe("none");
  });

  it("is not blocked with an annual plan even if team quantity is insufficient", () => {
    const annualOrgData = {
      ...trialExpiredOrgData,
      active_subscription_plan: "annual" as const,
    };
    renderWithUser(baseUser, annualOrgData);
    expect(screen.getByTestId("blocked").textContent).toBe("false");
  });
});

describe("useMaybeUpgradeBlock", () => {
  it("returns null when outside provider", () => {
    render(<MaybeTestConsumer />);
    expect(screen.getByTestId("has-context").textContent).toBe("false");
  });

  it("returns context when inside provider", () => {
    mockUseOrganization.mockReturnValue({
      data: compliantOrgData,
      isLoading: false,
      error: null,
    } as unknown as ReturnType<typeof useOrganization>);

    render(
      <QueryClientProvider client={queryClient}>
        <TestUserStoreProvider user={baseUser}>
          <UpgradeBlockProvider>
            <MaybeTestConsumer />
          </UpgradeBlockProvider>
        </TestUserStoreProvider>
      </QueryClientProvider>
    );
    expect(screen.getByTestId("has-context").textContent).toBe("true");
  });
});