js/tests/modelComplexity.test.js

/**
 * Unit tests for ModelComplexity.compute() and _formatParameterCount()
 * Validates: Requirements 20.1, 20.2, 20.3, 20.4, 20.5
 */

import { describe, it, expect } from 'vitest';

// Mirror the pure compute logic from ModelComplexity
function computeComplexity(parsedModel) {
  const nodes = (parsedModel && parsedModel.graph && parsedModel.graph.nodes) || [];
  const edges = (parsedModel && parsedModel.graph && parsedModel.graph.edges) || [];
  const initializers = (parsedModel && parsedModel.initializers) || [];

  let totalParameters = 0;
  let memoryFootprint = 0;

  for (const init of initializers) {
    totalParameters += (init.elementCount || 0);
    memoryFootprint += (init.size || 0);
  }

  return {
    totalParameters,
    memoryFootprint,
    totalNodes: nodes.length,
    totalEdges: edges.length,
    totalInitializers: initializers.length,
  };
}

// Mirror the formatting logic
function formatParameterCount(count) {
  if (count == null || isNaN(count) || count === 0) return '0';
  const abs = Math.abs(count);
  if (abs >= 1e9) return (count / 1e9).toFixed(1).replace(/\.0$/, '') + 'B';
  if (abs >= 1e6) return (count / 1e6).toFixed(1).replace(/\.0$/, '') + 'M';
  if (abs >= 1e3) return (count / 1e3).toFixed(1).replace(/\.0$/, '') + 'K';
  return String(count);
}

describe('ModelComplexity.compute', () => {
  it('should sum elementCount from all initializers as totalParameters', () => {
    const model = {
      graph: { nodes: [{ id: '1', opType: 'Conv' }], edges: [] },
      initializers: [
        { name: 'w1', elementCount: 1000, size: 4000 },
        { name: 'w2', elementCount: 500, size: 2000 },
        { name: 'b1', elementCount: 64, size: 256 },
      ],
    };
    const m = computeComplexity(model);
    expect(m.totalParameters).toBe(1564);
    expect(m.memoryFootprint).toBe(6256);
    expect(m.totalInitializers).toBe(3);
  });

  it('should sum size from all initializers as memoryFootprint', () => {
    const model = {
      graph: { nodes: [], edges: [] },
      initializers: [
        { name: 'w1', elementCount: 100, size: 400 },
        { name: 'w2', elementCount: 200, size: 800 },
      ],
    };
    const m = computeComplexity(model);
    expect(m.memoryFootprint).toBe(1200);
  });

  it('should count nodes, edges, and initializers', () => {
    const model = {
      graph: {
        nodes: [{ id: '1' }, { id: '2' }, { id: '3' }],
        edges: [{ source: '1', target: '2' }, { source: '2', target: '3' }],
      },
      initializers: [{ name: 'w1', elementCount: 10, size: 40 }],
    };
    const m = computeComplexity(model);
    expect(m.totalNodes).toBe(3);
    expect(m.totalEdges).toBe(2);
    expect(m.totalInitializers).toBe(1);
  });

  it('should handle no initializers (0 parameters)', () => {
    const model = {
      graph: { nodes: [{ id: '1' }], edges: [] },
      initializers: [],
    };
    const m = computeComplexity(model);
    expect(m.totalParameters).toBe(0);
    expect(m.memoryFootprint).toBe(0);
    expect(m.totalInitializers).toBe(0);
  });

  it('should handle null/undefined model', () => {
    const m = computeComplexity(null);
    expect(m.totalParameters).toBe(0);
    expect(m.memoryFootprint).toBe(0);
    expect(m.totalNodes).toBe(0);
    expect(m.totalEdges).toBe(0);
    expect(m.totalInitializers).toBe(0);
  });

  it('should handle initializers with missing elementCount or size', () => {
    const model = {
      graph: { nodes: [], edges: [] },
      initializers: [
        { name: 'w1' },
        { name: 'w2', elementCount: 100 },
        { name: 'w3', size: 500 },
      ],
    };
    const m = computeComplexity(model);
    expect(m.totalParameters).toBe(100);
    expect(m.memoryFootprint).toBe(500);
  });
});

describe('formatParameterCount', () => {
  it('should return "0" for zero', () => {
    expect(formatParameterCount(0)).toBe('0');
  });

  it('should return plain number for counts under 1000', () => {
    expect(formatParameterCount(500)).toBe('500');
    expect(formatParameterCount(999)).toBe('999');
  });

  it('should format thousands with K suffix', () => {
    expect(formatParameterCount(1500)).toBe('1.5K');
    expect(formatParameterCount(10000)).toBe('10K');
    expect(formatParameterCount(999999)).toBe('1000K');
  });

  it('should format millions with M suffix', () => {
    expect(formatParameterCount(1000000)).toBe('1M');
    expect(formatParameterCount(25600000)).toBe('25.6M');
    expect(formatParameterCount(138000000)).toBe('138M');
  });

  it('should format billions with B suffix', () => {
    expect(formatParameterCount(1200000000)).toBe('1.2B');
    expect(formatParameterCount(7000000000)).toBe('7B');
  });

  it('should return "0" for null/undefined/NaN', () => {
    expect(formatParameterCount(null)).toBe('0');
    expect(formatParameterCount(undefined)).toBe('0');
    expect(formatParameterCount(NaN)).toBe('0');
  });
});