[function] searchTree - find a node in a tree

searchTree.test.ts

import { searchTree } from '../searchTree';

const validTree = [
  {
    id: 1,
    parent_id: 0,
    title: 'Level 1',
    value: 1,
    children: [
      {
        id: 3,
        parent_id: 1,
        title: 'Level 1-1',
        value: 3,
        children: [
          {
            id: 5,
            parent_id: 3,
            title: 'Level 1-1-1',
            value: 5,
          },
          {
            id: 6,
            parent_id: 3,
            title: 'Level 1-1-2',
            value: 6,
          },
        ],
      },
      {
        id: 4,
        parent_id: 1,
        title: 'Level 1-2',
        value: 4,
      },
    ],
  },
  {
    id: 2,
    parent_id: 0,
    title: 'Level 2',
    value: 2,
    children: [
      {
        id: 7,
        parent_id: 2,
        title: 'Level 2-1',
        value: 7,
        children: [
          {
            id: 8,
            parent_id: 7,
            title: 'Level 2-1-1',
            value: 8,
          },
        ],
      },
    ],
  },
  {
    id: 9,
    parent_id: 0,
    title: 'Level 3',
    value: 9,
  },
];

describe('searchTree', () => {
  test('Invalid parameter should return null', () => {
    expect(searchTree(null as any, 'foo', 'bar')).toEqual(null);
    expect(searchTree([] as any, 'foo', 'bar')).toEqual(null);
    expect(searchTree({} as any, 'foo', 'bar')).toEqual(null);
    expect(searchTree(undefined as any, 'foo', 'bar')).toEqual(null);
    expect(searchTree(NaN as any, 'foo', 'bar')).toEqual(null);
    expect(searchTree(true as any, 'foo', 'bar')).toEqual(null);
    expect(searchTree(false as any, 'foo', 'bar')).toEqual(null);
    expect(searchTree('' as any, 'foo', 'bar')).toEqual(null);
  });
  test('Valid parameter should return correct object', () => {
    // level 3
    expect(searchTree(validTree, 5, 'value')).toEqual({
      id: 5,
      parent_id: 3,
      title: 'Level 1-1-1',
      value: 5,
    });
    expect(searchTree(validTree, 7, 'parent_id')).toEqual({
      id: 8,
      parent_id: 7,
      title: 'Level 2-1-1',
      value: 8,
    });
    // level 2
    expect(searchTree(validTree, 'Level 1-2', 'title')).toEqual({
      id: 4,
      parent_id: 1,
      title: 'Level 1-2',
      value: 4,
    });
    // level 1
    expect(searchTree(validTree, 9, 'id')).toEqual({
      id: 9,
      parent_id: 0,
      title: 'Level 3',
      value: 9,
    });
    // with children
    expect(searchTree(validTree, 3, 'id', true)).toEqual({
      id: 3,
      parent_id: 1,
      title: 'Level 1-1',
      value: 3,
      children: [
        {
          id: 5,
          parent_id: 3,
          title: 'Level 1-1-1',
          value: 5,
        },
        {
          id: 6,
          parent_id: 3,
          title: 'Level 1-1-2',
          value: 6,
        },
      ],
    });
  });
});

searchTree.ts

export function searchTree(
  tree: Record<string, any>[],
  value: unknown,
  key = 'value',
  withChildren = false,
) {
  let result = null;
  if (!Array.isArray(tree)) return result;

  for (let index = 0; index < tree.length; index += 1) {
    const stack = [tree[index]];
    while (stack.length) {
      const node = stack.shift()!;
      if (node[key] === value) {
        result = node;
        break;
      }
      if (node.children) {
        stack.push(...node.children);
      }
    }
    if (result) break;
  }
  if (withChildren !== true) {
    delete result?.children;
  }

  return result;
}