Bopytex/test/test_scheduler.py

from bopytex.tasks import Task
from bopytex.scheduler import Scheduler
import pytest


def action_done(deps, args, output):
    return f"{deps} - {args} - {output} - done"


actions = {"DO": action_done}


def test_schedule_append():
    scheduler = Scheduler(actions)
    tasks = [
        Task(action="DO", args={}, deps=["dep1", "dep2"], output="end1"),
        Task(action="DO", args={}, deps=["dep1", "dep3"], output="end2"),
    ]
    scheduler.append(tasks)
    assert scheduler.tasks == tasks
    assert scheduler.all_deps == {"dep1", "dep2", "dep3"}
    assert scheduler.all_output == {"end1", "end2"}


def test_schedule_dispatch():
    scheduler = Scheduler(actions)
    task = Task(action="DO", args={}, deps=[], output="end")
    result = scheduler.dispatch(task)
    assert result == "[] - {} - end - done"


def test_schedule_one_task():
    scheduler = Scheduler(actions)
    scheduler.append([Task(action="DO", args={}, deps=[], output="end")])
    result = scheduler.next()
    assert result == "[] - {} - end - done"
    assert scheduler.tasks == []
    assert scheduler.done == ["end"]


def test_schedule_one_task_with_args():
    scheduler = Scheduler(actions)
    scheduler.append([Task(action="DO", args={"task": "one"}, deps=[], output="one")])
    result = scheduler.next()
    assert result == "[] - {'task': 'one'} - one - done"
    assert scheduler.tasks == []
    assert scheduler.done == ["one"]


def test_schedule_multiple_tasks():
    scheduler = Scheduler(actions)
    t1 = Task(action="DO", args={"task": "one"}, deps=[], output="one")
    t2 = Task(action="DO", args={"task": "two"}, deps=[], output="two")
    t3 = Task(action="DO", args={"task": "three"}, deps=[], output="three")
    scheduler.append([t1, t2, t3])

    assert scheduler.is_finishable() == True

    result = scheduler.next()
    assert result == "[] - {'task': 'one'} - one - done"
    assert scheduler.tasks == [t2, t3]
    assert scheduler.done == ["one"]

    result = scheduler.next()
    assert result == "[] - {'task': 'two'} - two - done"
    assert scheduler.tasks == [t3]
    assert scheduler.done == ["one", "two"]

    result = scheduler.next()
    assert result == "[] - {'task': 'three'} - three - done"
    assert scheduler.tasks == []
    assert scheduler.done == ["one", "two", "three"]


def test_schedule_multiple_tasks_with_dependencies():
    scheduler = Scheduler(actions)
    t1 = Task(action="DO", args={"task": "one"}, deps=["three"], output="one")
    t2 = Task(action="DO", args={"task": "two"}, deps=["one"], output="two")
    t3 = Task(action="DO", args={"task": "three"}, deps=[], output="three")
    scheduler.append([t1, t2, t3])

    assert scheduler.is_finishable() == True

    result = scheduler.next()
    assert result == "[] - {'task': 'three'} - three - done"
    assert scheduler.tasks == [t1, t2]
    assert scheduler.done == ["three"]

    result = scheduler.next()
    assert result == "['three'] - {'task': 'one'} - one - done"
    assert scheduler.tasks == [t2]
    assert scheduler.done == ["three", "one"]

    result = scheduler.next()
    assert result == "['one'] - {'task': 'two'} - two - done"
    assert scheduler.tasks == []
    assert scheduler.done == ["three", "one", "two"]


def test_schedule_empty_task():
    scheduler = Scheduler(actions)
    scheduler.append([])
    with pytest.raises(StopIteration):
        scheduler.next()


def test_schedule_multiple_tasks_with_undoable_dependencies():
    scheduler = Scheduler(actions)
    t1 = Task(action="DO", args={"task": "one"}, deps=["three"], output="one")
    t2 = Task(action="DO", args={"task": "two"}, deps=[], output="two")
    scheduler.append([t1, t2])

    assert scheduler.is_finishable() == False

    scheduler.run()
    assert scheduler.tasks == [t1]
    assert scheduler.done == ["two"]