docmost-sync/test/collaboration-mutate.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import * as Y from 'yjs';

// -----------------------------------------------------------------------------
// Mock the Hocuspocus provider so the test drives its whole lifecycle by hand.
//
// `mutatePageContent` constructs `new HocuspocusProvider({...})` and wires all
// behaviour through callbacks on the OPTIONS object it passes in:
//   onConnect / onDisconnect / onClose / onSynced / onAuthenticationFailed.
// It does NOT use an `onUnsyncedChanges` option — instead it reads the live
// `provider.unsyncedChanges` getter and subscribes via
// `provider.on("unsyncedChanges", handler)` / `provider.off(...)`.
//
// The mock therefore:
//   - records EVERY constructed instance (so a test can assert "not called");
//   - captures the options object (to fire onSynced/onDisconnect/onClose/...);
//   - exposes the real `Y.Doc` passed as `options.document` — the source reads
//     and writes THAT doc (a closure variable), never `provider.document`, so
//     reading the fragment off the captured doc reflects what was written;
//   - implements `on`/`off` for the "unsyncedChanges" event, a settable
//     `unsyncedChanges` getter-backed field, and no-op `destroy`/`disconnect`.
// -----------------------------------------------------------------------------

type UnsyncedHandler = (data: { number: number }) => void;

interface MockProviderHandle {
  opts: any;
  ydoc: Y.Doc;
  unsyncedChanges: number;
  listeners: Map<string, Set<(...args: any[]) => void>>;
  destroyed: boolean;
  // Helpers to drive the lifecycle from a test.
  fireSynced: () => void;
  fireDisconnect: () => void;
  fireClose: () => void;
  fireAuthFailed: () => void;
  fireConnect: () => void;
  // Set unsyncedChanges and emit the "unsyncedChanges" event with the new value.
  emitUnsynced: (n: number) => void;
}

// Shared registry the test reads after invoking the SUT.
const instances: MockProviderHandle[] = [];

vi.mock('@hocuspocus/provider', () => {
  const HocuspocusProvider = vi.fn().mockImplementation((opts: any) => {
    const listeners = new Map<string, Set<(...args: any[]) => void>>();
    const handle: MockProviderHandle = {
      opts,
      ydoc: opts.document as Y.Doc,
      unsyncedChanges: 1, // default: write is outstanding until told otherwise
      listeners,
      destroyed: false,
      fireSynced: () => opts.onSynced && opts.onSynced(),
      fireDisconnect: () => opts.onDisconnect && opts.onDisconnect(),
      fireClose: () => opts.onClose && opts.onClose(),
      fireAuthFailed: () =>
        opts.onAuthenticationFailed && opts.onAuthenticationFailed(),
      fireConnect: () => opts.onConnect && opts.onConnect(),
      emitUnsynced: (n: number) => {
        handle.unsyncedChanges = n;
        const set = listeners.get('unsyncedChanges');
        if (set) for (const fn of set) fn({ number: n });
      },
    };

    // The object the SUT actually interacts with.
    const provider: any = {
      // `unsyncedChanges` is read synchronously by waitForPersistence; back it
      // by the handle so a test can preset it before firing onSynced.
      get unsyncedChanges() {
        return handle.unsyncedChanges;
      },
      on: (event: string, fn: (...args: any[]) => void) => {
        if (!listeners.has(event)) listeners.set(event, new Set());
        listeners.get(event)!.add(fn);
      },
      off: (event: string, fn: (...args: any[]) => void) => {
        listeners.get(event)?.delete(fn);
      },
      destroy: () => {
        handle.destroyed = true;
      },
      disconnect: () => {},
      document: opts.document,
    };

    // Let a test reach the provider object too, if needed.
    (handle as any).provider = provider;
    instances.push(handle);
    return provider;
  });
  return { HocuspocusProvider };
});

// Import AFTER vi.mock so the mocked provider is in place. Import directly from
// the source .js (matches the repo's other tests, e.g. page-lock.test.ts).
import {
  mutatePageContent,
  replacePageContent,
} from '../packages/docmost-client/src/lib/collaboration.js';
import { HocuspocusProvider } from '@hocuspocus/provider';

// A valid minimal ProseMirror doc used as the "new" content to write.
function newDocWith(text: string) {
  return {
    type: 'doc',
    content: [
      {
        type: 'paragraph',
        content: [{ type: 'text', text }],
      },
    ],
  };
}

// Read the "default" XML fragment off a Y.Doc and report whether it has nodes.
function fragmentLength(ydoc: Y.Doc): number {
  return ydoc.getXmlFragment('default').length;
}

// Flush pending microtasks so callbacks fired synchronously settle before the
// test inspects results. withPageLock chains through several microtask hops.
async function flushMicrotasks() {
  for (let i = 0; i < 20; i++) await Promise.resolve();
}

// Drive the SUT: call mutatePageContent, wait for the provider to be built,
// run `drive(handle)` to fire lifecycle callbacks, then await the result.
async function runMutate(
  transform: (live: any) => any,
  drive: (handle: MockProviderHandle) => void | Promise<void>,
  pageId = uniquePageId(),
) {
  const promise = mutatePageContent(pageId, 'collab-token', 'http://x/api', transform);
  // The provider is constructed only AFTER the page lock grants the turn,
  // which is a few microtask hops in. Wait until the instance shows up.
  await flushMicrotasks();
  const handle = instances[instances.length - 1];
  expect(handle, 'provider should have been constructed').toBeTruthy();
  await drive(handle);
  await flushMicrotasks();
  return promise;
}

// Unique pageId per test: the page lock is keyed by a process-global Map, so a
// reused id would serialize unrelated tests behind each other.
let pageCounter = 0;
function uniquePageId() {
  return `page-${process.pid}-${Date.now()}-${pageCounter++}`;
}

beforeEach(() => {
  vi.useFakeTimers();
  instances.length = 0;
});

afterEach(() => {
  vi.useRealTimers();
  vi.clearAllMocks();
});

describe('replacePageContent — fail-fast guard (before any provider)', () => {
  it('throws on null document WITHOUT constructing a provider', async () => {
    await expect(
      replacePageContent('p1', null as any, 'tok', 'http://x/api'),
    ).rejects.toThrow(/invalid ProseMirror document/);
    // The guard runs before mutatePageContent, so no provider is built.
    expect(HocuspocusProvider).not.toHaveBeenCalled();
    expect(instances.length).toBe(0);
  });

  it('throws on a non-"doc"-typed object WITHOUT constructing a provider', async () => {
    await expect(
      replacePageContent(
        'p2',
        { type: 'paragraph', content: [] } as any,
        'tok',
        'http://x/api',
      ),
    ).rejects.toThrow(/invalid ProseMirror document/);
    expect(HocuspocusProvider).not.toHaveBeenCalled();
    expect(instances.length).toBe(0);
  });

  it('throws on a non-object (string) WITHOUT constructing a provider', async () => {
    await expect(
      replacePageContent('p3', 'just a string' as any, 'tok', 'http://x/api'),
    ).rejects.toThrow(/invalid ProseMirror document/);
    expect(HocuspocusProvider).not.toHaveBeenCalled();
    expect(instances.length).toBe(0);
  });
});

describe('mutatePageContent — read/transform/write core', () => {
  it('transform receives the default empty doc when the live Y.Doc is empty', async () => {
    let received: any;
    const promise = await runMutate(
      (live) => {
        received = live;
        return null; // abort: we only care about what the transform saw
      },
      (h) => {
        // Empty live doc -> onSynced reads it and runs the transform.
        h.fireSynced();
      },
    );
    await promise;
    // An empty fragment yields the synthesized default doc.
    expect(received).toEqual({ type: 'doc', content: [] });
  });

  it('transform returning null performs NO write and resolves with the live doc', async () => {
    let liveDocRef: any;
    const promise = await runMutate(
      (live) => {
        liveDocRef = live;
        return null;
      },
      (h) => h.fireSynced(),
    );
    const result = await promise;
    // No write: the captured Y.Doc fragment is still empty.
    const handle = instances[instances.length - 1];
    expect(fragmentLength(handle.ydoc)).toBe(0);
    // The returned value is the live doc (the default empty doc here).
    expect(result).toBe(liveDocRef);
    expect(result).toEqual({ type: 'doc', content: [] });
  });

  it('transform throwing propagates: the returned promise rejects', async () => {
    const pageId = uniquePageId();
    const promise = mutatePageContent(pageId, 'tok', 'http://x/api', () => {
      throw new Error('boom from transform');
    });
    // Attach the rejection handler up-front so the rejection is observed even
    // though the throw is surfaced synchronously inside the SUT's onSynced
    // try/catch (which converts it into a finish()/reject()).
    const settled = promise.then(
      () => ({ ok: true as const }),
      (err: Error) => ({ ok: false as const, err }),
    );
    await flushMicrotasks();
    const handle = instances[instances.length - 1];
    handle.fireSynced();
    await flushMicrotasks();
    const outcome = await settled;
    expect(outcome.ok).toBe(false);
    if (!outcome.ok) {
      expect(outcome.err.message).toMatch(/boom from transform/);
    }
  });

  it('transform returning a new doc replaces the fragment (old gone, new present)', async () => {
    const pageId = uniquePageId();
    const promise = mutatePageContent(
      pageId,
      'tok',
      'http://x/api',
      () => newDocWith('hello world'),
    );
    await flushMicrotasks();
    const handle = instances[instances.length - 1];
    expect(handle).toBeTruthy();

    // Seed the live doc with some pre-existing content so we can prove it is
    // fully replaced (not merged). We write a paragraph into the fragment that
    // the SUT will read on sync, then clear+rewrite.
    const liveFragment = handle.ydoc.getXmlFragment('default');
    handle.ydoc.transact(() => {
      const el = new Y.XmlElement('paragraph');
      el.insert(0, [new Y.XmlText('OLD CONTENT')]);
      liveFragment.insert(0, [el]);
    });
    expect(fragmentLength(handle.ydoc)).toBe(1);

    // Fire sync: the SUT reads liveDoc, runs transform, deletes the fragment,
    // and applies the new doc's update.
    handle.fireSynced();
    // The write is synchronous; persistence is still pending (unsyncedChanges=1).
    await flushMicrotasks();

    // The fragment now reflects the NEW doc. The old "OLD CONTENT" text is gone.
    const xml = handle.ydoc.getXmlFragment('default').toString();
    expect(xml).toContain('hello world');
    expect(xml).not.toContain('OLD CONTENT');

    // Now acknowledge persistence so the promise resolves cleanly.
    handle.emitUnsynced(0);
    await flushMicrotasks();
    const result = await promise;
    expect(result).toEqual(newDocWith('hello world'));
  });
});

describe('mutatePageContent — persistence / false-success suppression', () => {
  it('unsyncedChanges->0 while connected RESOLVES with the written doc', async () => {
    const pageId = uniquePageId();
    const promise = mutatePageContent(
      pageId,
      'tok',
      'http://x/api',
      () => newDocWith('persisted'),
    );
    await flushMicrotasks();
    const handle = instances[instances.length - 1];
    // Outstanding write at sync time (unsyncedChanges defaults to 1), so the
    // SUT subscribes and waits for the event.
    handle.fireSynced();
    await flushMicrotasks();

    // Server acknowledges: counter drops to 0 while still connected.
    handle.emitUnsynced(0);
    await flushMicrotasks();
    const result = await promise;
    expect(result).toEqual(newDocWith('persisted'));
  });

  it('resolves immediately when unsyncedChanges is already 0 at persist-check time', async () => {
    const pageId = uniquePageId();
    const promise = mutatePageContent(
      pageId,
      'tok',
      'http://x/api',
      () => newDocWith('already synced'),
    );
    await flushMicrotasks();
    const handle = instances[instances.length - 1];
    // Pretend the write was acknowledged before waitForPersistence checks.
    handle.unsyncedChanges = 0;
    handle.fireSynced();
    await flushMicrotasks();
    const result = await promise;
    expect(result).toEqual(newDocWith('already synced'));
  });

  it('disconnect BEFORE reaching 0 unsynced does NOT resolve as success (connectionLost guard)', async () => {
    const pageId = uniquePageId();
    const promise = mutatePageContent(
      pageId,
      'tok',
      'http://x/api',
      () => newDocWith('should not persist'),
    );
    // Reject handler attached up-front so the rejection is never unhandled.
    const settled = promise.then(
      () => ({ ok: true as const }),
      (err: Error) => ({ ok: false as const, err }),
    );
    await flushMicrotasks();
    const handle = instances[instances.length - 1];

    // Sync + write happen; persistence subscription is registered (unsynced=1).
    handle.fireSynced();
    await flushMicrotasks();

    // Connection drops: this sets connectionLost=true and finishes with an error.
    handle.fireDisconnect();
    await flushMicrotasks();

    // A late reconnect handshake drives the counter back to 0. Because the
    // connection was already lost, the unsyncedChanges handler must NOT report
    // success — and finish() is idempotent (settled flag), so this is a no-op.
    handle.emitUnsynced(0);
    await flushMicrotasks();

    const outcome = await settled;
    expect(outcome.ok).toBe(false);
    if (!outcome.ok) {
      expect(outcome.err.message).toMatch(
        /connection closed before the update was persisted/i,
      );
    }
  });

  it('onClose before persistence rejects with the connection-closed error', async () => {
    const pageId = uniquePageId();
    const promise = mutatePageContent(
      pageId,
      'tok',
      'http://x/api',
      () => newDocWith('x'),
    );
    const settled = promise.then(
      () => ({ ok: true as const }),
      (err: Error) => ({ ok: false as const, err }),
    );
    await flushMicrotasks();
    const handle = instances[instances.length - 1];
    handle.fireSynced();
    await flushMicrotasks();
    handle.fireClose();
    await flushMicrotasks();
    const outcome = await settled;
    expect(outcome.ok).toBe(false);
    if (!outcome.ok) {
      expect(outcome.err.message).toMatch(
        /connection closed before the update was persisted/i,
      );
    }
  });

  it('authentication failure rejects with the auth error', async () => {
    const pageId = uniquePageId();
    const promise = mutatePageContent(
      pageId,
      'tok',
      'http://x/api',
      () => newDocWith('x'),
    );
    const settled = promise.then(
      () => ({ ok: true as const }),
      (err: Error) => ({ ok: false as const, err }),
    );
    await flushMicrotasks();
    const handle = instances[instances.length - 1];
    handle.fireAuthFailed();
    await flushMicrotasks();
    const outcome = await settled;
    expect(outcome.ok).toBe(false);
    if (!outcome.ok) {
      expect(outcome.err.message).toMatch(/Authentication failed/i);
    }
  });
});

describe('mutatePageContent — timeouts (fake timers)', () => {
  it('connect timeout rejects when onSynced never fires', async () => {
    const pageId = uniquePageId();
    const promise = mutatePageContent(
      pageId,
      'tok',
      'http://x/api',
      () => newDocWith('never'),
    );
    const settled = promise.then(
      () => ({ ok: true as const }),
      (err: Error) => ({ ok: false as const, err }),
    );
    await flushMicrotasks();
    // Never fire onSynced. Advance past CONNECT_TIMEOUT_MS (25000).
    await vi.advanceTimersByTimeAsync(25001);
    await flushMicrotasks();
    const outcome = await settled;
    expect(outcome.ok).toBe(false);
    if (!outcome.ok) {
      expect(outcome.err.message).toMatch(/Connection timeout/i);
    }
  });

  it('persist timeout rejects when unsyncedChanges never reaches 0 after the write', async () => {
    const pageId = uniquePageId();
    const promise = mutatePageContent(
      pageId,
      'tok',
      'http://x/api',
      () => newDocWith('stuck'),
    );
    const settled = promise.then(
      () => ({ ok: true as const }),
      (err: Error) => ({ ok: false as const, err }),
    );
    await flushMicrotasks();
    const handle = instances[instances.length - 1];
    // Write happens, but the counter stays at 1 (default) — never acknowledged.
    handle.fireSynced();
    await flushMicrotasks();
    // Advance past PERSIST_TIMEOUT_MS (20000). The connect timer was cleared
    // when onSynced ran, so only the persist timer is pending.
    await vi.advanceTimersByTimeAsync(20001);
    await flushMicrotasks();
    const outcome = await settled;
    expect(outcome.ok).toBe(false);
    if (!outcome.ok) {
      expect(outcome.err.message).toMatch(/persist the update/i);
    }
  });
});