site-builder/craft/src/utils/craft-tree.test.ts

import { describe, test, expect } from 'vitest';
import type { NodeTree, Node } from '@craftjs/core';
import type { SerializedTreeNode } from '../types/sitesmith';
import { regenerateTreeIds, flattenTreeForCraft, sanitizeAiTree } from './craft-tree';

function makeNode(overrides: Partial<Node>): Node {
  return {
    id: 'placeholder',
    data: {
      props: {},
      type: 'div',
      name: 'div',
      displayName: 'div',
      isCanvas: false,
      parent: null,
      linkedNodes: {},
      nodes: [],
      hidden: false,
    },
    info: {},
    events: { selected: false, dragged: false, hovered: false },
    dom: null,
    related: {},
    rules: {
      canDrag: () => true,
      canDrop: () => true,
      canMoveIn: () => true,
      canMoveOut: () => true,
    },
    _hydrationTimestamp: 0,
    ...overrides,
  } as Node;
}

function makeTree(): NodeTree {
  // root -> [childA, childB]; root also has a linkedNodes entry -> childB
  const root = makeNode({
    id: 'root-1',
    data: {
      ...makeNode({}).data,
      parent: null,
      nodes: ['child-a-1'],
      linkedNodes: { slot: 'child-b-1' },
    },
  });
  const childA = makeNode({
    id: 'child-a-1',
    data: {
      ...makeNode({}).data,
      parent: 'root-1',
      nodes: [],
      linkedNodes: {},
    },
  });
  const childB = makeNode({
    id: 'child-b-1',
    data: {
      ...makeNode({}).data,
      parent: 'root-1',
      nodes: [],
      linkedNodes: {},
    },
  });

  return {
    rootNodeId: 'root-1',
    nodes: {
      'root-1': root,
      'child-a-1': childA,
      'child-b-1': childB,
    },
  };
}

describe('regenerateTreeIds', () => {
  test('produces an id set fully disjoint from the input', () => {
    const input = makeTree();
    const inputIds = Object.keys(input.nodes);
    const output = regenerateTreeIds(input);
    const outputIds = Object.keys(output.nodes);

    for (const id of outputIds) {
      expect(inputIds).not.toContain(id);
    }
  });

  test('rootNodeId is a key in nodes and equals that node id', () => {
    const output = regenerateTreeIds(makeTree());
    expect(output.nodes[output.rootNodeId]).toBeDefined();
    expect(output.nodes[output.rootNodeId].id).toBe(output.rootNodeId);
  });

  test('every child data.parent points to a NEW id that exists in the output', () => {
    const output = regenerateTreeIds(makeTree());
    for (const id of Object.keys(output.nodes)) {
      const node = output.nodes[id];
      if (node.data.parent !== null) {
        expect(output.nodes[node.data.parent]).toBeDefined();
      }
    }
    // specifically the two children under the (new) root
    const newRoot = output.nodes[output.rootNodeId];
    expect(newRoot.data.nodes.length).toBe(1);
    const newChildAId = newRoot.data.nodes[0];
    expect(output.nodes[newChildAId].data.parent).toBe(output.rootNodeId);
  });

  test('linkedNodes values are remapped to existing new ids', () => {
    const output = regenerateTreeIds(makeTree());
    const newRoot = output.nodes[output.rootNodeId];
    const linkedId = newRoot.data.linkedNodes.slot;
    expect(linkedId).toBeDefined();
    expect(output.nodes[linkedId]).toBeDefined();
    // linked node's parent should still reference the new root
    expect(output.nodes[linkedId].data.parent).toBe(output.rootNodeId);
  });

  test('structure/count is preserved', () => {
    const input = makeTree();
    const output = regenerateTreeIds(input);
    expect(Object.keys(output.nodes).length).toBe(Object.keys(input.nodes).length);
    expect(output.nodes[output.rootNodeId].data.nodes.length).toBe(
      input.nodes[input.rootNodeId].data.nodes.length
    );
  });

  test('does not mutate the input tree', () => {
    const input = makeTree();
    const snapshot = JSON.parse(JSON.stringify(input));
    regenerateTreeIds(input);
    expect(JSON.parse(JSON.stringify(input))).toEqual(snapshot);
  });

  test('output node data.props is a deep clone, not shared with the input node', () => {
    const input = makeTree();
    input.nodes['root-1'].data.props = { alignment: 'left' };

    const output = regenerateTreeIds(input);
    const outputRoot = output.nodes[output.rootNodeId];

    // Clone must preserve content at the time of cloning.
    expect(outputRoot.data.props).toEqual({ alignment: 'left' });

    // Mutating the output's props (simulating Craft.js's in-place setProp)
    // must NOT affect the input node's props object.
    (outputRoot.data.props as { alignment: string }).alignment = 'right';

    expect(input.nodes['root-1'].data.props).toEqual({ alignment: 'left' });
  });
});

describe('flattenTreeForCraft', () => {
  test('style: [] normalizes to {}', () => {
    const tree: SerializedTreeNode = {
      type: { resolvedName: 'Heading' },
      props: { node_id: 'h1', style: [] as unknown as Record<string, unknown> },
      nodes: [],
    };
    const { nodes } = flattenTreeForCraft(tree);
    expect(nodes['h1'].props.style).toEqual({});
  });

  test('a ColumnLayout with direct children flattens them into col-N linkedNodes', () => {
    const tree: SerializedTreeNode = {
      type: { resolvedName: 'ColumnLayout' },
      props: { node_id: 'cols' },
      nodes: [
        { type: { resolvedName: 'Heading' }, props: { node_id: 'a' }, nodes: [] },
        { type: { resolvedName: 'Heading' }, props: { node_id: 'b' }, nodes: [] },
      ],
    };
    const { nodes } = flattenTreeForCraft(tree);
    expect(nodes['cols'].nodes).toEqual([]);
    expect(nodes['cols'].linkedNodes).toEqual({ 'col-0': 'a', 'col-1': 'b' });
    expect(nodes['cols'].props.columns).toBe(2);
    expect(nodes['a'].parent).toBe('cols');
    expect(nodes['a'].isCanvas).toBe(true);
    expect(nodes['b'].parent).toBe('cols');
  });

  test('a Section with direct children wraps them in a synthetic section-inner Container', () => {
    const tree: SerializedTreeNode = {
      type: { resolvedName: 'Section' },
      props: { node_id: 'sec' },
      nodes: [
        { type: { resolvedName: 'Heading' }, props: { node_id: 'h1' }, nodes: [] },
      ],
    };
    const { rootNodeId, nodes, syntheticIds } = flattenTreeForCraft(tree);
    expect(rootNodeId).toBe('sec');
    const innerId = nodes['sec'].linkedNodes['section-inner'];
    expect(innerId).toBe('sec__section-inner');
    expect(syntheticIds.has(innerId)).toBe(true);
    expect(nodes[innerId].nodes).toEqual(['h1']);
    expect(nodes[innerId].type.resolvedName).toBe('Container');
    expect(nodes['h1'].parent).toBe(innerId);
    expect(nodes['sec'].nodes).toEqual([]);
  });

  test('a small tree with both a ColumnLayout and a SHELL_INNER wrapper flattens the same way regardless of caller intent', () => {
    // Same shape both PageContext.treeToCraftState and apply-ai-response's
    // buildNodeTree feed through this helper — the walk itself must not care
    // which one is calling it.
    const tree: SerializedTreeNode = {
      type: { resolvedName: 'Section' },
      props: { node_id: 'sec' },
      nodes: [
        {
          type: { resolvedName: 'ColumnLayout' },
          props: { node_id: 'cols' },
          nodes: [
            { type: { resolvedName: 'Heading' }, props: { node_id: 'a' }, nodes: [] },
          ],
        },
      ],
    };
    const run1 = flattenTreeForCraft(tree);
    const run2 = flattenTreeForCraft(tree);
    expect(run1.rootNodeId).toBe(run2.rootNodeId);
    expect(Object.keys(run1.nodes).sort()).toEqual(Object.keys(run2.nodes).sort());
    const innerId = run1.nodes['sec'].linkedNodes['section-inner'];
    expect(run1.nodes[innerId].linkedNodes).toEqual({});
    expect(run1.nodes[innerId].nodes).toEqual(['cols']);
    expect(run1.nodes['cols'].linkedNodes).toEqual({ 'col-0': 'a' });
    expect(run1.nodes['a'].parent).toBe('cols');
  });

  test('auto-generates an id when node_id is missing', () => {
    const tree: SerializedTreeNode = { type: { resolvedName: 'Heading' }, props: {}, nodes: [] };
    const { rootNodeId, nodes } = flattenTreeForCraft(tree);
    expect(typeof rootNodeId).toBe('string');
    expect(nodes[rootNodeId]).toBeDefined();
  });

  test('isCanvas is true only for Container', () => {
    const containerTree: SerializedTreeNode = { type: { resolvedName: 'Container' }, props: { node_id: 'c1' }, nodes: [] };
    const headingTree: SerializedTreeNode = { type: { resolvedName: 'Heading' }, props: { node_id: 'h1' }, nodes: [] };
    expect(flattenTreeForCraft(containerTree).nodes['c1'].isCanvas).toBe(true);
    expect(flattenTreeForCraft(headingTree).nodes['h1'].isCanvas).toBe(false);
  });
});

describe('sanitizeAiTree (re-exported for treeToCraftState parity with buildNodeTree)', () => {
  test('unknown resolvedName at the root is rejected (returns null)', () => {
    const tree = { type: { resolvedName: 'NotARealComponent' }, props: {}, nodes: [] };
    expect(sanitizeAiTree(tree, new Set())).toBeNull();
  });

  test('unknown resolvedName on a nested child drops only that subtree', () => {
    const tree = {
      type: { resolvedName: 'Section' },
      props: { node_id: 's1' },
      nodes: [
        { type: { resolvedName: 'NotARealComponent' }, props: {}, nodes: [] },
        { type: { resolvedName: 'Heading' }, props: { node_id: 'h1' }, nodes: [] },
      ],
    };
    const out = sanitizeAiTree(tree, new Set());
    expect(out!.nodes!.length).toBe(1);
    expect(out!.nodes![0].props.node_id).toBe('h1');
  });
});