Repository: dask/dask
File: dask/tests/test_task_spec.py (Python)

Code before:
def test_fuse_tasks_key():
    a = Task("key-1", func, "a", "b")
    b = Task("key-2", func2, a.ref(), "d")
    for t1, t2 in itertools.permutations((a, b)):
        fused = Task.fuse(t2, t1)
        assert fused.key == b.key

        fused = Task.fuse(t2, t1, key="new-key")
        assert fused.key == "new-key"


def test_fuse_tasks():
    a = Task("key-1", func, "a", "b")
    b = Task("key-2", func2, a.ref(), "d")
    for t1, t2 in itertools.permutations((a, b)):
        fused = Task.fuse(t2, t1)

        assert fused() == func2(func("a", "b"), "d")

        t1 = Task("key-1", func, TaskRef("dependency"), "b")
        t2 = Task("key-2", func2, t1.ref(), "d")

        fused = Task.fuse(t2, t1)
        assert fused.dependencies == {"dependency"}

        assert fused({"dependency": "dep"}) == func2(func("dep", "b"), "d")


def test_fuse_reject_multiple_outputs():
    a = Task("key-1", func, "a", "b")
    b = Task("key-2", func2, "a", "d")
    for t1, t2 in itertools.permutations((a, b)):
        with pytest.raises(ValueError, match="multiple outputs"):
            Task.fuse(t1, t2)


def test_fused_ensure_only_executed_once():
    counter = []

    def counter_func(a, b):
        counter.append(None)
        return func(a, b)

    a = Task("key-1", counter_func, "a", "a")
    b = Task("key-2", func2, a.ref(), "b")
    c = Task("key-3", func2, a.ref(), "c")
    d = Task("key-4", func, b.ref(), c.ref())
    for perm in itertools.permutations([a, b, c, d]):
        fused = Task.fuse(*perm)
        counter.clear()
        assert fused() == func(func2(func("a", "a"), "b"), func2(func("a", "a"), "c"))

Review comment (suggestion):
```suggestion
def test_fuse_tasks():
    a = Task("key-1", func, "a", "b")
    b = Task("key-2", func2, a.ref(), "d")
    c = Task("key-3", func3, b.ref(), "e")
    for t1, t2, t3 in itertools.permutations((a, b, c)):
        fused = Task.fuse(t3, t2, t1)

        assert fused() == func3(func2(func("a", "b"), "d"), "e")

        t1 = Task("key-1", func, TaskRef("dependency"), "b")
        t2 = Task("key-2", func2, t1.ref(), "d")
        t3 = Task("key-3", func3, t2.ref(), "e")
        fused = Task.fuse(t3, t2, t1)
        assert fused.dependencies == {"dependency"}

        assert fused({"dependency": "dep"}) == func3(func2(func("dep", "b"), "d"), "e")
```

It took me a moment to understand that fuse is indeed order-independent, even in the case of transitive dependencies. Extending this test may or may not help. Feel free to ignore.

Code after:
def test_fuse_tasks_key():
    a = Task("key-1", func, "a", "b")
    b = Task("key-2", func2, a.ref(), "d")
    for t1, t2 in itertools.permutations((a, b)):
        fused = Task.fuse(t2, t1)
        assert fused.key == b.key

        fused = Task.fuse(t2, t1, key="new-key")
        assert fused.key == "new-key"


def test_fuse_tasks():
    a = Task("key-1", func, "a", "b")
    b = Task("key-2", func2, a.ref(), "d")
    c = Task("key-3", func3, b.ref(), "e")
    for t1, t2, t3 in itertools.permutations((a, b, c)):
        fused = Task.fuse(t3, t2, t1)
        assert fused() == func3(func2(func("a", "b"), "d"), "e")
        t1 = Task("key-1", func, TaskRef("dependency"), "b")
        t2 = Task("key-2", func2, t1.ref(), "d")
        t3 = Task("key-3", func3, t2.ref(), "e")
        fused = Task.fuse(t3, t2, t1)
        assert fused.dependencies == {"dependency"}
        assert fused({"dependency": "dep"}) == func3(func2(func("dep", "b"), "d"), "e")


def test_fuse_reject_multiple_outputs():
    a = Task("key-1", func, "a", "b")
    b = Task("key-2", func2, "a", "d")
    for t1, t2 in itertools.permutations((a, b)):
        with pytest.raises(ValueError, match="multiple outputs"):
            Task.fuse(t1, t2)


def test_fused_ensure_only_executed_once():
    counter = []

    def counter_func(a, b):
        counter.append(None)
        return func(a, b)

    a = Task("key-1", counter_func, "a", "a")
    b = Task("key-2", func2, a.ref(), "b")
    c = Task("key-3", func2, a.ref(), "c")
    d = Task("key-4", func, b.ref(), c.ref())
    for perm in itertools.permutations([a, b, c, d]):
        fused = Task.fuse(*perm)
        counter.clear()
        assert fused() == func(func2(func("a", "a"), "b"), func2(func("a", "a"), "c"))
        assert len(counter) == 1