Shade/packages/shade-streams/tests/partition.test.ts

import { describe, test, expect } from 'bun:test';
import { ValidationError } from '@shade/core';
import {
  planRangePartition,
  planRoundRobinPartition,
  chunkRange,
  partitionsEqual,
} from '../src/index.js';

describe('planRangePartition', () => {
  test('evenly divisible totalBytes', () => {
    const lanes = planRangePartition(100, 4);
    expect(lanes).toHaveLength(4);
    expect(lanes[0]!.partition).toEqual({ kind: 'range', startByte: 0, endByte: 25, startChunk: 0 });
    expect(lanes[3]!.partition).toEqual({ kind: 'range', startByte: 75, endByte: 100, startChunk: 0 });
  });

  test('non-divisible: extra bytes go to early lanes', () => {
    const lanes = planRangePartition(10, 3);
    const ranges = lanes.map((l) => l.partition);
    // 10 / 3 = 3 remainder 1 — lane 0 gets 4, lanes 1+2 get 3 each
    expect(ranges[0]).toEqual({ kind: 'range', startByte: 0, endByte: 4, startChunk: 0 });
    expect(ranges[1]).toEqual({ kind: 'range', startByte: 4, endByte: 7, startChunk: 0 });
    expect(ranges[2]).toEqual({ kind: 'range', startByte: 7, endByte: 10, startChunk: 0 });
  });

  test('lanes cover entire range without gaps or overlap', () => {
    for (const total of [0, 1, 7, 100, 1024 * 1024]) {
      for (const count of [1, 2, 4, 16]) {
        const lanes = planRangePartition(total, count);
        let cursor = 0;
        for (const lane of lanes) {
          if (lane.partition.kind !== 'range') throw new Error('expected range');
          expect(lane.partition.startByte).toBe(cursor);
          cursor = lane.partition.endByte;
        }
        expect(cursor).toBe(total);
      }
    }
  });

  test('1-lane partition spans entire input', () => {
    const lanes = planRangePartition(500, 1);
    expect(lanes).toHaveLength(1);
    expect(lanes[0]!.partition).toEqual({ kind: 'range', startByte: 0, endByte: 500, startChunk: 0 });
  });

  test('rejects negative totalBytes / fractional / non-positive count', () => {
    expect(() => planRangePartition(-1, 1)).toThrow(ValidationError);
    expect(() => planRangePartition(1.5, 1)).toThrow(ValidationError);
    expect(() => planRangePartition(100, 0)).toThrow(ValidationError);
    expect(() => planRangePartition(100, -1)).toThrow(ValidationError);
  });

  test('laneIds are 0..count-1 in order', () => {
    const lanes = planRangePartition(64, 16);
    for (let i = 0; i < 16; i++) expect(lanes[i]!.laneId).toBe(i);
  });
});

describe('planRoundRobinPartition', () => {
  test('produces N lanes labeled 0..N-1', () => {
    const lanes = planRoundRobinPartition(8);
    expect(lanes).toHaveLength(8);
    for (let i = 0; i < 8; i++) {
      expect(lanes[i]!.laneId).toBe(i);
      expect(lanes[i]!.partition).toEqual({ kind: 'round-robin', lane: i, count: 8 });
    }
  });

  test('rejects non-positive count', () => {
    expect(() => planRoundRobinPartition(0)).toThrow(ValidationError);
    expect(() => planRoundRobinPartition(-1)).toThrow(ValidationError);
  });
});

describe('chunkRange', () => {
  test('splits an even range into chunkSize slices', () => {
    expect(chunkRange(0, 1024, 256)).toEqual([
      { start: 0, end: 256 },
      { start: 256, end: 512 },
      { start: 512, end: 768 },
      { start: 768, end: 1024 },
    ]);
  });

  test('last chunk truncated for non-divisible range', () => {
    expect(chunkRange(0, 1000, 256)).toEqual([
      { start: 0, end: 256 },
      { start: 256, end: 512 },
      { start: 512, end: 768 },
      { start: 768, end: 1000 },
    ]);
  });

  test('non-zero start offset is preserved', () => {
    expect(chunkRange(100, 350, 100)).toEqual([
      { start: 100, end: 200 },
      { start: 200, end: 300 },
      { start: 300, end: 350 },
    ]);
  });

  test('empty range produces a single empty chunk (so isLast can be carried)', () => {
    expect(chunkRange(50, 50, 100)).toEqual([{ start: 50, end: 50 }]);
  });

  test('rejects non-positive chunkSize', () => {
    expect(() => chunkRange(0, 100, 0)).toThrow(ValidationError);
    expect(() => chunkRange(0, 100, -1)).toThrow(ValidationError);
  });
});

describe('partitionsEqual', () => {
  test('identical range partitions', () => {
    expect(
      partitionsEqual(
        { kind: 'range', startByte: 0, endByte: 100, startChunk: 0 },
        { kind: 'range', startByte: 0, endByte: 100, startChunk: 0 },
      ),
    ).toBe(true);
  });

  test('different range bounds', () => {
    expect(
      partitionsEqual(
        { kind: 'range', startByte: 0, endByte: 100, startChunk: 0 },
        { kind: 'range', startByte: 0, endByte: 200, startChunk: 0 },
      ),
    ).toBe(false);
  });

  test('range vs round-robin → false', () => {
    expect(
      partitionsEqual(
        { kind: 'range', startByte: 0, endByte: 100, startChunk: 0 },
        { kind: 'round-robin', lane: 0, count: 1 },
      ),
    ).toBe(false);
  });

  test('identical round-robin partitions', () => {
    expect(
      partitionsEqual(
        { kind: 'round-robin', lane: 2, count: 4 },
        { kind: 'round-robin', lane: 2, count: 4 },
      ),
    ).toBe(true);
  });

  test('different round-robin lane index', () => {
    expect(
      partitionsEqual(
        { kind: 'round-robin', lane: 1, count: 4 },
        { kind: 'round-robin', lane: 2, count: 4 },
      ),
    ).toBe(false);
  });
});