def test__DAG_add_dependencies(): A, B = get_two_tasks() C = Task("C.py", env="test-env") dag = DAG() dag.add_dependencies({B: A}) assert dag._edges[A] == set([B]) dag = DAG() dag.add_dependencies({C: {A, B}}) assert dag._edges[A] == set([C]) assert dag._edges[B] == set([C])
def test__DAG_add_task(): A, B = get_two_tasks() dag = DAG() dag.add_task(A) assert dag.tasks == {A}, "Test Task was not added to the DAG"
def test__DAG_get_upstream(): A, B = get_two_tasks() dag = DAG() dag.add_dependency(B, depends_on=A) assert dag.get_upstream() is not None assert dag.get_upstream()[B] == {A} assert dag.get_upstream() == {B: {A}}, "Task A is not upstream"
def test__DAG_get_downstream(): A, B = get_two_tasks() dag = DAG() dag.add_dependency(B, depends_on=A) assert dag.get_downstream() is not None assert dag.get_downstream()[A] == {B} assert dag.get_downstream() == {A: {B}}, "Task B is not downstream"
def test__DAG_add_dependency_detect_cycle2(): A, B = get_two_tasks() C = Task("C.py", env="test-env") dag = DAG() with pytest.raises(CyclicGraphError): dag.add_dependencies({A: C, B: A, C: B})
def test__Composer_get_schedules(): A = Task("A.py", "test-exe") B = Task("B.py", "test-exe") C = Task("C.py", "test-exe") Z = Task("Z.py", "test-exe") dag = DAG() dag.add_tasks({A, B, C, Z}) dag.add_dependencies({B: A, C: B}) dq = Composer(dag) priorities = dq.get_schedules() testable = {} # build a testable result dict for k, v in priorities.items(): new_set = set() if isinstance(v, Task): testable[k] = set(v) continue for vi in v: new_set.add(hash(vi)) testable[k] = new_set assert testable == {1: {hash(A), hash(Z)}, 2: {hash(B)}, 3: {hash(C)}}
def test__DAG_remove_task(): A, B = get_two_tasks() dag = DAG() dag.add_tasks({A, B}) dag.remove_task(A) assert dag.tasks == {B}
def test__DAG_init(): DAG() # init with dependencies make_tea = Task("make_tea.py", "test-env") drink_tea = Task("drink_tea.py", "test-env") dag = DAG(tasks=make_tea) assert len(dag.tasks) == 1 dag = DAG(tasks={drink_tea, make_tea}) assert len(dag.tasks) == 2 dag = DAG(upstream_dependencies={drink_tea: make_tea}) assert len(dag.tasks) == 2 dag = DAG(downstream_dependencies={make_tea: drink_tea}) assert len(dag.tasks) == 2
def test__DAG_is_cyclic(): A, B = get_two_tasks() dag = DAG() dag.add_dependency(B, depends_on=A) assert not dag.is_cyclic(), "acyclic graph idenfied as cyclic" with pytest.raises(CyclicGraphError): dag.add_dependency(A, depends_on=B)
def test__Composer_repr(): p = pathlib.Path("make_tea.py") make_tea = Task(p, "test-exe") dag = DAG() dag.add_task(make_tea) dq = Composer(dag) assert repr(dq) == "".join( ["Composer(DAG({Task(", p.resolve().as_posix(), ")}))"])
def test__DAG_remove_tasks(): A, B = get_two_tasks() C = Task("C.py") dag = DAG() dag.add_tasks({A, B, C}) dag.remove_tasks({A, B}) assert dag.tasks == {C} dag.remove_tasks(C) assert dag.tasks == set()
def test__DAG_add_tasks(): A, B = get_two_tasks() C = Task("C.py") dag = DAG() dag.add_tasks({A, B}) assert dag.tasks == {A, B}, "Test Tasks were not added to the DAG" dag.add_tasks(C) assert dag.tasks == {A, B, C}
def test__Composer_init(): """Nothing should break here """ A = Task("A.py", "test-exe") B = Task("B.py", "test-exe") C = Task("C.py", "test-exe") dag = DAG() dag.add_tasks({A, B, C}) dq = Composer(dag) return None
def test__Composer_init_exceptions(): """Raise expected exceptions """ A = Task("A.py", "test-exe") B = Task("B.py", "test-exe") C = Task("C.py", "test-exe") dag = DAG() dag.add_tasks({A, B, C}) with pytest.raises(TypeError): Composer() return None
def test__Composer_get_task_schedules(): A = Task("A.py", "test-exe") B = Task("B.py", "test-exe") C = Task("C.py", "test-exe") Z = Task("Z.py", "test-exe") dag = DAG() dag.add_tasks({A, B, C, Z}) dag.add_dependencies({B: A, C: B}) dq = Composer(dag) priorities = dq.get_task_schedules() testable = {hash(k): v for k, v in priorities.items()} assert testable == {hash(A): 1, hash(B): 2, hash(C): 3, hash(Z): 1}
def test__readme_example(): ## define tasks and environments pour_water = Task("./tea-tasks/pour_water.py") boil_water = Task("./tea-tasks/boil_water.py") prep_infuser = Task("./tea-tasks/prep_infuser.py") steep_tea = Task("./tea-tasks/steep_tea.py") ## define runtime dependencies make_tea = DAG(upstream_dependencies={ boil_water: {pour_water}, steep_tea: {boil_water, prep_infuser}, }) ## run tasks dq = Composer(make_tea) dq.get_schedules()
def test__Composer_refresh_dag(): A = Task("A.py", "test-exe") B = Task("B.py", "test-exe") C = Task("C.py", "test-exe") dag = DAG() dag.add_tasks({A, B, C}) dq = Composer(dag) tasks = sorted(list(dq.dag.tasks)) for t in tasks: dq.dag.remove_task(t) assert dq.dag.tasks == set() dq.refresh_dag() new_tasks = sorted(list(dq.dag.tasks)) for t, nt in zip(tasks, new_tasks): assert t == nt
def test__DAG_get_sinks(): A, B = get_two_tasks() dag = DAG() dag.add_dependency(B, depends_on=A) assert dag.get_sinks() is not None assert dag.get_sinks() == {B}
def test__DAG_add_dependency_detect_cycle(): A, B = get_two_tasks() dag = DAG() dag.add_dependency(B, A) with pytest.raises(CyclicGraphError): dag.add_dependency(A, B)
def test__DAG_add_dependency(): A, B = get_two_tasks() dag = DAG() dag.add_dependency(B, A) assert dag._edges[A] == set([B])
def test__DAG_repr(): p = pathlib.Path("make_tea.py") make_tea = Task(p, "test-env") dag = DAG() dag.add_task(make_tea) assert repr(dag) == "".join(["DAG({Task(", p.resolve().as_posix(), ")})"])