CedarCopilot/cedar-mail

import React, { createRef } from 'react';
import { act, render, waitFor } from '@testing-library/react';

import { Document, type DocumentHandle } from '@/modules/documents/document';

const mockGetDoc = jest.fn();
const mockUseDocEvents = jest.fn();
const mockAcquireProvider = jest.fn();
const mockReleaseProvider = jest.fn();
const mockGetProvider = jest.fn();
const mockBase64ToUint8Array = jest.fn((...args: unknown[]) => new Uint8Array([String(args[0]).length]));

jest.mock('@/providers/query-provider', () => ({
  trpcClient: {
    documents: {
      getDoc: {
        query: (...args: unknown[]) => mockGetDoc(...args),
      },
    },
  },
  useTRPCClient: () => ({
    documents: {
      getDoc: {
        query: (...args: unknown[]) => mockGetDoc(...args),
      },
    },
  }),
}));

jest.mock('@/modules/documents/yjs', () => ({
  EnsureNodeIds: {},
  createTrpcApplyUpdateClient: jest.fn(() => ({ applyUpdate: jest.fn() })),
  acquireProvider: (...args: unknown[]) => mockAcquireProvider(...args),
  releaseProvider: (...args: unknown[]) => mockReleaseProvider(...args),
  getProvider: (...args: unknown[]) => mockGetProvider(...args),
  base64ToUint8Array: (...args: unknown[]) => mockBase64ToUint8Array(...args),
  useDocEvents: (...args: unknown[]) => mockUseDocEvents(...args),
}));

jest.mock('@/components/markdown-editor', () => {
  const ReactModule = require('react');
  const MarkdownEditor = ReactModule.forwardRef(function MockMarkdownEditor(
    props: {
      onEditorCreated?: (editor: unknown) => void;
      children?: React.ReactNode;
    },
    ref: React.Ref<unknown>,
  ) {
    ReactModule.useImperativeHandle(ref, () => ({
      editor: { id: 'editor-1' },
      getMarkdown: () => 'mock-markdown',
    }));
    ReactModule.useEffect(() => {
      props.onEditorCreated?.({ id: 'editor-1' });
    }, [props.onEditorCreated]);
    return <div data-testid="mock-markdown-editor">{props.children}</div>;
  });

  return {
    MarkdownEditor,
  };
});

type MockProvider = {
  idbReady: Promise<void>;
  ydoc: { getXmlFragment: jest.Mock };
  applyInitialState: jest.Mock;
  forceFlush: jest.Mock<Promise<void>, []>;
};

function makeProvider(fragmentLength: number): MockProvider {
  return {
    idbReady: Promise.resolve(),
    ydoc: {
      getXmlFragment: jest.fn(() => ({ length: fragmentLength })),
    },
    applyInitialState: jest.fn(),
    forceFlush: jest.fn(async () => {}),
  };
}

describe('<Document /> lifecycle', () => {
  beforeEach(() => {
    jest.clearAllMocks();
  });

  it('always seeds from server, even when IDB is warm', async () => {
    // A warm IDB only proves the client has every edit it has ever
    // observed — agents/other tabs may have written while we weren't
    // watching, so we must merge in the server state on every mount.
    const provider = makeProvider(1);
    mockAcquireProvider.mockReturnValue(provider);
    mockGetProvider.mockReturnValue(provider);
    mockGetDoc.mockResolvedValue({ contentYjs: 'warm-seed' });

    render(<Document documentId="doc-warm" />);

    await waitFor(() => expect(mockGetDoc).toHaveBeenCalledWith({ documentId: 'doc-warm' }));
    await waitFor(() => expect(mockBase64ToUint8Array).toHaveBeenCalledWith('warm-seed'));
    expect(provider.applyInitialState).toHaveBeenCalledWith(new Uint8Array([9]));
    expect(mockUseDocEvents).toHaveBeenCalledWith('doc-warm');
  });

  it('fetches and seeds content on cold open', async () => {
    const provider = makeProvider(0);
    mockAcquireProvider.mockReturnValue(provider);
    mockGetProvider.mockReturnValue(provider);
    mockGetDoc.mockResolvedValue({ contentYjs: 'seed-bytes' });

    render(<Document documentId="doc-cold" />);

    await waitFor(() => expect(mockGetDoc).toHaveBeenCalledWith({ documentId: 'doc-cold' }));
    await waitFor(() => expect(mockBase64ToUint8Array).toHaveBeenCalledWith('seed-bytes'));
    expect(provider.applyInitialState).toHaveBeenCalledWith(new Uint8Array([10]));
  });

  it('forwards targetUserId to getDoc so an admin viewing a teammate seeds their copy', async () => {
    // Regression: AgentInstructionsSection passed targetUserId to itself but never
    // down into Document, so an org admin viewing a teammate's agent got a getDoc
    // call scoped to the admin's OWN session — the server rejected it as a
    // cross-user read and the editor rendered permanently blank.
    const provider = makeProvider(1);
    mockAcquireProvider.mockReturnValue(provider);
    mockGetProvider.mockReturnValue(provider);
    mockGetDoc.mockResolvedValue({ contentYjs: 'teammate-seed' });

    render(<Document documentId="doc-teammate" targetUserId="teammate-1" />);

    await waitFor(() =>
      expect(mockGetDoc).toHaveBeenCalledWith({
        documentId: 'doc-teammate',
        targetUserId: 'teammate-1',
      }),
    );
  });

  it('refreshFromServer force-flushes and re-seeds the provider', async () => {
    const provider = makeProvider(1);
    mockAcquireProvider.mockReturnValue(provider);
    mockGetProvider.mockReturnValue(provider);
    mockGetDoc.mockResolvedValue({ contentYjs: 'fresh-seed' });

    const ref = createRef<DocumentHandle>();
    render(<Document ref={ref} documentId="doc-refresh" />);

    await waitFor(() => expect(ref.current?.documentId).toBe('doc-refresh'));

    await act(async () => {
      await ref.current?.refreshFromServer();
    });

    expect(provider.forceFlush).toHaveBeenCalledTimes(1);
    expect(mockGetDoc).toHaveBeenCalledWith({ documentId: 'doc-refresh' });
    expect(mockBase64ToUint8Array).toHaveBeenCalledWith('fresh-seed');
    expect(provider.applyInitialState).toHaveBeenCalledWith(new Uint8Array([10]));
  });

  it('refreshFromServer also forwards targetUserId', async () => {
    const provider = makeProvider(1);
    mockAcquireProvider.mockReturnValue(provider);
    mockGetProvider.mockReturnValue(provider);
    mockGetDoc.mockResolvedValue({ contentYjs: 'fresh-teammate-seed' });

    const ref = createRef<DocumentHandle>();
    render(<Document ref={ref} documentId="doc-refresh-teammate" targetUserId="teammate-1" />);

    await waitFor(() => expect(ref.current?.documentId).toBe('doc-refresh-teammate'));

    await act(async () => {
      await ref.current?.refreshFromServer();
    });

    expect(mockGetDoc).toHaveBeenCalledWith({
      documentId: 'doc-refresh-teammate',
      targetUserId: 'teammate-1',
    });
  });
});