from eval.distributed_eval_utils import distributed_batch_plan


def test_distributed_batch_plan_keeps_uneven_ranks_in_final_sync() -> None:
    plan = distributed_batch_plan(total_items=2500, num_processes=8, batch_size=1)

    assert plan.local_item_counts == (313, 313, 313, 313, 312, 312, 312, 312)
    assert plan.sync_batches == 313


def test_distributed_batch_plan_handles_batched_uneven_shards() -> None:
    plan = distributed_batch_plan(total_items=10, num_processes=3, batch_size=2)

    assert plan.local_item_counts == (4, 3, 3)
    assert plan.sync_batches == 2