from truecluster.cluster.planner import plan_even_layers


def test_even_plan():
    assignments = plan_even_layers(num_layers=24, max_nodes=3)
    assert [(a.layer_start, a.layer_end_exclusive) for a in assignments] == [(0, 8), (8, 16), (16, 24)]


def test_remainder_plan():
    assignments = plan_even_layers(num_layers=25, max_nodes=4)
    assert [(a.layer_start, a.layer_end_exclusive) for a in assignments] == [(0, 7), (7, 13), (13, 19), (19, 25)]