From 56eab51e55cec75048d5d0c18742eb864eb09e06 Mon Sep 17 00:00:00 2001
From: Piotr Wysocki <hello@piotrwysocki.com>
Date: Thu, 16 May 2024 21:09:44 +0200
Subject: [PATCH] Make graph model independent

---
 smartschedule/parallelization/stage.py        | 23 -----
 .../parallelization/stage_parallelization.py  | 27 ------
 smartschedule/planning/__init__.py            |  4 +
 .../parallelization/__init__.py               |  0
 .../parallelization/parallel_stages.py        |  2 +-
 .../parallelization/parallel_stages_list.py   |  8 +-
 .../sorted_nodes_to_parallelized_stages.py    | 17 ++++
 .../planning/parallelization/stage.py         | 23 +++++
 .../parallelization/stage_parallelization.py  | 15 ++++
 .../parallelization/stages_to_nodes.py        | 44 ++++++++++
 smartschedule/sorter/__init__.py              |  6 ++
 .../sorter/feedback_arc_set_on_graph.py       | 45 ++++++++++
 .../sorter/graph_topological_sort.py          | 22 +++++
 smartschedule/sorter/node.py                  | 17 ++++
 smartschedule/sorter/nodes.py                 | 17 ++++
 smartschedule/sorter/sorted_nodes.py          | 26 ++++++
 tests/planning/__init__.py                    |  0
 tests/planning/parallelization/__init__.py    |  0
 .../parallelization/test_parallelization.py   | 54 ++++++++++++
 tests/sorter/__init__.py                      |  0
 .../sorter/test_feedback_arc_set_on_graph.py  | 35 ++++++++
 tests/sorter/test_graph_topological_sort.py   | 84 +++++++++++++++++++
 tests/test_parallelization.py                 | 36 --------
 23 files changed, 413 insertions(+), 92 deletions(-)
 delete mode 100644 smartschedule/parallelization/stage.py
 delete mode 100644 smartschedule/parallelization/stage_parallelization.py
 create mode 100644 smartschedule/planning/__init__.py
 rename smartschedule/{ => planning}/parallelization/__init__.py (100%)
 rename smartschedule/{ => planning}/parallelization/parallel_stages.py (74%)
 rename smartschedule/{ => planning}/parallelization/parallel_stages_list.py (60%)
 create mode 100644 smartschedule/planning/parallelization/sorted_nodes_to_parallelized_stages.py
 create mode 100644 smartschedule/planning/parallelization/stage.py
 create mode 100644 smartschedule/planning/parallelization/stage_parallelization.py
 create mode 100644 smartschedule/planning/parallelization/stages_to_nodes.py
 create mode 100644 smartschedule/sorter/__init__.py
 create mode 100644 smartschedule/sorter/feedback_arc_set_on_graph.py
 create mode 100644 smartschedule/sorter/graph_topological_sort.py
 create mode 100644 smartschedule/sorter/node.py
 create mode 100644 smartschedule/sorter/nodes.py
 create mode 100644 smartschedule/sorter/sorted_nodes.py
 create mode 100644 tests/planning/__init__.py
 create mode 100644 tests/planning/parallelization/__init__.py
 create mode 100644 tests/planning/parallelization/test_parallelization.py
 create mode 100644 tests/sorter/__init__.py
 create mode 100644 tests/sorter/test_feedback_arc_set_on_graph.py
 create mode 100644 tests/sorter/test_graph_topological_sort.py
 delete mode 100644 tests/test_parallelization.py

diff --git a/smartschedule/parallelization/stage.py b/smartschedule/parallelization/stage.py
deleted file mode 100644
index 918c12d..0000000
--- a/smartschedule/parallelization/stage.py
+++ /dev/null
@@ -1,23 +0,0 @@
-from datetime import timedelta
-from attrs import frozen, field
-
-
-@frozen
-class ResourceName:
-    name: str
-
-
-@frozen
-class Stage:
-    stage_name: str
-    dependencies: set["Stage"] = field(factory=set, hash=False)
-    resources: set[ResourceName] = field(factory=set, hash=False)
-    duration: timedelta = field(default=timedelta.min, hash=False)
-
-    def depends_on(self, stage: "Stage") -> "Stage":
-        self.dependencies.add(stage)
-        return self
-
-    @property
-    def name(self):
-        return self.stage_name
diff --git a/smartschedule/parallelization/stage_parallelization.py b/smartschedule/parallelization/stage_parallelization.py
deleted file mode 100644
index 0fd0411..0000000
--- a/smartschedule/parallelization/stage_parallelization.py
+++ /dev/null
@@ -1,27 +0,0 @@
-from graphlib import TopologicalSorter, CycleError
-
-from smartschedule.parallelization.parallel_stages import ParallelStages
-from smartschedule.parallelization.parallel_stages_list import ParallelStagesList
-from smartschedule.parallelization.stage import Stage
-
-
-class StageParallelization:
-    @staticmethod
-    def of(stages: set[Stage]) -> ParallelStagesList:
-        parallel_stages_list = ParallelStagesList()
-        sorter = TopologicalSorter(
-            graph={stage: stage.dependencies for stage in stages}
-        )
-
-        try:
-            sorter.prepare()
-        except CycleError:
-            return parallel_stages_list
-
-        while sorter.is_active():
-            group = sorter.get_ready()
-            parallel_stages = ParallelStages(set(group))
-            parallel_stages_list = parallel_stages_list.add(parallel_stages)
-            sorter.done(*group)
-
-        return parallel_stages_list
diff --git a/smartschedule/planning/__init__.py b/smartschedule/planning/__init__.py
new file mode 100644
index 0000000..293bf2b
--- /dev/null
+++ b/smartschedule/planning/__init__.py
@@ -0,0 +1,4 @@
+__all__ = ["ResourceName", "Stage", "stage_parallelization_of"]
+
+from .parallelization.stage import ResourceName, Stage
+from .parallelization.stage_parallelization import stage_parallelization_of
diff --git a/smartschedule/parallelization/__init__.py b/smartschedule/planning/parallelization/__init__.py
similarity index 100%
rename from smartschedule/parallelization/__init__.py
rename to smartschedule/planning/parallelization/__init__.py
diff --git a/smartschedule/parallelization/parallel_stages.py b/smartschedule/planning/parallelization/parallel_stages.py
similarity index 74%
rename from smartschedule/parallelization/parallel_stages.py
rename to smartschedule/planning/parallelization/parallel_stages.py
index 4c16a65..5dcd027 100644
--- a/smartschedule/parallelization/parallel_stages.py
+++ b/smartschedule/planning/parallelization/parallel_stages.py
@@ -1,6 +1,6 @@
 from attrs import frozen
 
-from smartschedule.parallelization.stage import Stage
+from smartschedule.planning.parallelization.stage import Stage
 
 
 @frozen
diff --git a/smartschedule/parallelization/parallel_stages_list.py b/smartschedule/planning/parallelization/parallel_stages_list.py
similarity index 60%
rename from smartschedule/parallelization/parallel_stages_list.py
rename to smartschedule/planning/parallelization/parallel_stages_list.py
index 16d0d7f..62e20e9 100644
--- a/smartschedule/parallelization/parallel_stages_list.py
+++ b/smartschedule/planning/parallelization/parallel_stages_list.py
@@ -1,8 +1,6 @@
-from typing import Self
+from attrs import field, frozen
 
-from attrs import frozen, field
-
-from smartschedule.parallelization.parallel_stages import ParallelStages
+from smartschedule.planning.parallelization.parallel_stages import ParallelStages
 
 
 @frozen
@@ -12,6 +10,6 @@ class ParallelStagesList:
     def __str__(self) -> str:
         return " | ".join(str(parallel_stages) for parallel_stages in self.all)
 
-    def add(self, new_parallel_stages: ParallelStages) -> Self:
+    def add(self, new_parallel_stages: ParallelStages) -> "ParallelStagesList":
         result = [*self.all, new_parallel_stages]
         return ParallelStagesList(all=result)
diff --git a/smartschedule/planning/parallelization/sorted_nodes_to_parallelized_stages.py b/smartschedule/planning/parallelization/sorted_nodes_to_parallelized_stages.py
new file mode 100644
index 0000000..b7f7867
--- /dev/null
+++ b/smartschedule/planning/parallelization/sorted_nodes_to_parallelized_stages.py
@@ -0,0 +1,17 @@
+from smartschedule.planning.parallelization.parallel_stages import ParallelStages
+from smartschedule.planning.parallelization.parallel_stages_list import (
+    ParallelStagesList,
+)
+from smartschedule.planning.parallelization.stage import Stage
+from smartschedule.sorter import SortedNodes
+
+
+def sorted_nodes_to_parallelized_stages(
+    sorted_nodes: SortedNodes[Stage],
+) -> ParallelStagesList:
+    return ParallelStagesList(
+        [
+            ParallelStages({node.content for node in nodes if node.content is not None})
+            for nodes in sorted_nodes
+        ]
+    )
diff --git a/smartschedule/planning/parallelization/stage.py b/smartschedule/planning/parallelization/stage.py
new file mode 100644
index 0000000..919ea72
--- /dev/null
+++ b/smartschedule/planning/parallelization/stage.py
@@ -0,0 +1,23 @@
+from datetime import timedelta
+from typing import Self
+
+from attrs import evolve, field, frozen
+
+
+@frozen
+class ResourceName:
+    name: str
+
+
+@frozen
+class Stage:
+    name: str
+    dependencies: set["Stage"] = field(factory=set, eq=False)
+    resources: set[ResourceName] = field(factory=set, eq=False)
+    duration: timedelta = field(default=timedelta.min, eq=False)
+
+    def depends_on(self, stage: Self) -> Self:
+        return evolve(self, dependencies={*self.dependencies, stage})
+
+    def with_chosen_resource_capabilities(self, *resources: ResourceName) -> Self:
+        return evolve(self, resources={*self.resources, *resources})
diff --git a/smartschedule/planning/parallelization/stage_parallelization.py b/smartschedule/planning/parallelization/stage_parallelization.py
new file mode 100644
index 0000000..70970e8
--- /dev/null
+++ b/smartschedule/planning/parallelization/stage_parallelization.py
@@ -0,0 +1,15 @@
+from smartschedule.planning.parallelization.parallel_stages_list import (
+    ParallelStagesList,
+)
+from smartschedule.planning.parallelization.sorted_nodes_to_parallelized_stages import (
+    sorted_nodes_to_parallelized_stages,
+)
+from smartschedule.planning.parallelization.stage import Stage
+from smartschedule.planning.parallelization.stages_to_nodes import stages_to_nodes
+from smartschedule.sorter import graph_topological_sort
+
+
+def stage_parallelization_of(stages: set[Stage]) -> ParallelStagesList:
+    nodes = stages_to_nodes(stages)
+    sorted_nodes = graph_topological_sort(nodes)
+    return sorted_nodes_to_parallelized_stages(sorted_nodes)
diff --git a/smartschedule/planning/parallelization/stages_to_nodes.py b/smartschedule/planning/parallelization/stages_to_nodes.py
new file mode 100644
index 0000000..83be616
--- /dev/null
+++ b/smartschedule/planning/parallelization/stages_to_nodes.py
@@ -0,0 +1,44 @@
+from collections.abc import Collection
+from itertools import islice
+
+from smartschedule.planning.parallelization.stage import Stage
+from smartschedule.sorter import Node, Nodes
+
+type StageName = str
+type NodesByName = dict[StageName, Node[Stage]]
+
+
+def stages_to_nodes(stages: Collection[Stage]) -> Nodes[Stage]:
+    nodes = {stage.name: Node(name=stage.name, content=stage) for stage in stages}
+
+    for n, stage in enumerate(stages):
+        _explicit_dependencies(stage, nodes)
+        _shared_resources(stage, list(islice(stages, n + 1, None)), nodes)
+
+    return Nodes(nodes.values())
+
+
+def _explicit_dependencies(stage: Stage, nodes: NodesByName) -> NodesByName:
+    for other in stage.dependencies:
+        nodes[stage.name] = nodes[stage.name].depends_on(nodes[other.name])
+
+    return nodes
+
+
+def _shared_resources(
+    stage: Stage, with_stages: Collection[Stage], nodes: NodesByName
+) -> NodesByName:
+    for other in with_stages:
+        if stage.name == other.name:
+            continue
+
+        # No shared resources.
+        if stage.resources.isdisjoint(other.resources):
+            continue
+
+        if len(other.resources) > len(stage.resources):
+            nodes[stage.name] = nodes[stage.name].depends_on(nodes[other.name])
+        else:
+            nodes[other.name] = nodes[other.name].depends_on(nodes[stage.name])
+
+    return nodes
diff --git a/smartschedule/sorter/__init__.py b/smartschedule/sorter/__init__.py
new file mode 100644
index 0000000..370eb4e
--- /dev/null
+++ b/smartschedule/sorter/__init__.py
@@ -0,0 +1,6 @@
+__all__ = ["graph_topological_sort", "Node", "Nodes", "SortedNodes"]
+
+from .graph_topological_sort import graph_topological_sort
+from .node import Node
+from .nodes import Nodes
+from .sorted_nodes import SortedNodes
diff --git a/smartschedule/sorter/feedback_arc_set_on_graph.py b/smartschedule/sorter/feedback_arc_set_on_graph.py
new file mode 100644
index 0000000..fe0e3f8
--- /dev/null
+++ b/smartschedule/sorter/feedback_arc_set_on_graph.py
@@ -0,0 +1,45 @@
+from collections import defaultdict
+
+from attrs import frozen
+
+from smartschedule.sorter.node import Node
+
+type AdjacencyList = defaultdict[int, list[int]]
+
+
+@frozen
+class Edge:
+    source: int
+    target: int
+
+    def __str__(self) -> str:
+        return f"({self.source} -> {self.target})"
+
+
+def calculate_feedback_arc_set_on_graph(initial_nodes: list[Node[str]]) -> list[Edge]:
+    adjacency_list = _create_adjacency_list(initial_nodes)
+    feedback_edges: list[Edge] = []
+    visited: list[int] = [0] * (len(adjacency_list) + 1)
+
+    for i in adjacency_list:
+        neighbours = adjacency_list[i]
+        visited[i] = 1
+        for neighbour in neighbours:
+            if visited[neighbour] == 1:
+                feedback_edges.append(Edge(i, neighbour))
+            else:
+                visited[neighbour] = 1
+
+    return feedback_edges
+
+
+def _create_adjacency_list(initial_nodes: list[Node[str]]) -> AdjacencyList:
+    adjacency_list: AdjacencyList = defaultdict(list)
+
+    for i, node in enumerate(initial_nodes):
+        dependencies = [
+            initial_nodes.index(dependency) + 1 for dependency in node.dependencies
+        ]
+        adjacency_list[i + 1] = dependencies
+
+    return adjacency_list
diff --git a/smartschedule/sorter/graph_topological_sort.py b/smartschedule/sorter/graph_topological_sort.py
new file mode 100644
index 0000000..14d5ea7
--- /dev/null
+++ b/smartschedule/sorter/graph_topological_sort.py
@@ -0,0 +1,22 @@
+from graphlib import CycleError, TopologicalSorter
+
+from smartschedule.sorter.nodes import Nodes
+from smartschedule.sorter.sorted_nodes import SortedNodes
+
+
+def graph_topological_sort[T](nodes: Nodes[T]) -> SortedNodes[T]:
+    sorted_nodes = SortedNodes.empty()
+
+    sorter = TopologicalSorter(graph={node: node.dependencies for node in nodes})
+
+    try:
+        sorter.prepare()
+    except CycleError:
+        return sorted_nodes
+
+    while sorter.is_active():
+        group = sorter.get_ready()
+        sorted_nodes = sorted_nodes.add(Nodes(group))
+        sorter.done(*group)
+
+    return sorted_nodes
diff --git a/smartschedule/sorter/node.py b/smartschedule/sorter/node.py
new file mode 100644
index 0000000..45b80d5
--- /dev/null
+++ b/smartschedule/sorter/node.py
@@ -0,0 +1,17 @@
+from itertools import chain
+from typing import Self
+
+from attrs import field, frozen
+
+
+@frozen
+class Node[T]:
+    name: str
+    dependencies: set["Node[T]"] = field(factory=set, eq=False)
+    content: T | None = field(eq=False, default=None)
+
+    def depends_on(self, node: Self) -> "Node[T]":
+        return Node(self.name, set(chain(self.dependencies, {node})), self.content)
+
+    def __str__(self) -> str:
+        return self.name
diff --git a/smartschedule/sorter/nodes.py b/smartschedule/sorter/nodes.py
new file mode 100644
index 0000000..482340f
--- /dev/null
+++ b/smartschedule/sorter/nodes.py
@@ -0,0 +1,17 @@
+from typing import Collection
+
+from smartschedule.sorter.node import Node
+
+
+class Nodes[T]:
+    def __init__(self, nodes: Collection[Node[T]]) -> None:
+        self._nodes = nodes if isinstance(nodes, set) else set(nodes)
+
+    def __iter__(self):
+        return iter(self._nodes)
+
+    def __contains__(self, item):
+        return item in self._nodes
+
+    def __len__(self):
+        return len(self._nodes)
diff --git a/smartschedule/sorter/sorted_nodes.py b/smartschedule/sorter/sorted_nodes.py
new file mode 100644
index 0000000..0881652
--- /dev/null
+++ b/smartschedule/sorter/sorted_nodes.py
@@ -0,0 +1,26 @@
+from typing import Self
+
+from attrs import evolve, field, frozen
+
+from smartschedule.sorter.nodes import Nodes
+
+
+@frozen
+class SortedNodes[T]:
+    _groups: list[Nodes[T]] = field(converter=lambda x: list(x))
+
+    def add(self, nodes: Nodes[T], /) -> Self:
+        return evolve(self, groups=[*self._groups, nodes])
+
+    def __iter__(self):
+        return iter(self._groups)
+
+    def __len__(self) -> int:
+        return len(self._groups)
+
+    def __getitem__(self, item):
+        return self._groups[item]
+
+    @classmethod
+    def empty(cls) -> Self:
+        return cls([])
diff --git a/tests/planning/__init__.py b/tests/planning/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/planning/parallelization/__init__.py b/tests/planning/parallelization/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/planning/parallelization/test_parallelization.py b/tests/planning/parallelization/test_parallelization.py
new file mode 100644
index 0000000..d7ca699
--- /dev/null
+++ b/tests/planning/parallelization/test_parallelization.py
@@ -0,0 +1,54 @@
+from smartschedule.planning import ResourceName, Stage, stage_parallelization_of
+
+LEON = ResourceName("Leon")
+SLAWEK = ResourceName("Sławek")
+ERYK = ResourceName("Eric")
+KUBA = ResourceName("Kuba")
+
+
+def test_everything_can_be_done_in_parallel_when_there_are_no_deps() -> None:
+    stage_1 = Stage("stage_1")
+    stage_2 = Stage("stage_2")
+
+    sorted_stages = stage_parallelization_of({stage_1, stage_2})
+
+    assert len(sorted_stages.all) == 1
+
+
+def test_simple_deps() -> None:
+    stage_1 = Stage("stage_1")
+    stage_2 = Stage("stage_2")
+    stage_3 = Stage("stage_3")
+    stage_4 = Stage("stage_4")
+    stage_2 = stage_2.depends_on(stage_1)
+    stage_3 = stage_3.depends_on(stage_1)
+    stage_4 = stage_4.depends_on(stage_2)
+
+    sorted_stages = stage_parallelization_of({stage_1, stage_2, stage_3, stage_4})
+
+    assert str(sorted_stages) == "stage_1 | stage_2, stage_3 | stage_4"
+
+
+def test_cannot_be_done_when_there_is_a_cycle() -> None:
+    stage_1 = Stage("stage_1")
+    stage_2 = Stage("stage_2")
+    stage_2 = stage_2.depends_on(stage_1)
+    stage_1 = stage_1.depends_on(stage_2)  # making it cyclic
+
+    sorted_stages = stage_parallelization_of({stage_1, stage_2})
+
+    assert len(sorted_stages.all) == 0
+
+
+def test_takes_into_account_shared_resources() -> None:
+    stage_1 = Stage("stage_1").with_chosen_resource_capabilities(LEON)
+    stage_2 = Stage("stage_2").with_chosen_resource_capabilities(ERYK, LEON)
+    stage_3 = Stage("stage_3").with_chosen_resource_capabilities(SLAWEK)
+    stage_4 = Stage("stage_4").with_chosen_resource_capabilities(SLAWEK, KUBA)
+
+    parallel_stages = stage_parallelization_of({stage_1, stage_2, stage_3, stage_4})
+
+    assert str(parallel_stages) in [
+        "stage_1, stage_3 | stage_2, stage_4",
+        "stage_2, stage_4 | stage_1, stage_3",
+    ]
diff --git a/tests/sorter/__init__.py b/tests/sorter/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/sorter/test_feedback_arc_set_on_graph.py b/tests/sorter/test_feedback_arc_set_on_graph.py
new file mode 100644
index 0000000..ec9d114
--- /dev/null
+++ b/tests/sorter/test_feedback_arc_set_on_graph.py
@@ -0,0 +1,35 @@
+from smartschedule.sorter.feedback_arc_set_on_graph import (
+    Edge,
+    calculate_feedback_arc_set_on_graph,
+)
+from smartschedule.sorter.node import Node
+
+
+def test_can_find_minimum_number_of_edges_to_remove_to_make_the_graph_acyclic() -> None:
+    node1 = Node("1")
+    node2 = Node("2")
+    node3 = Node("3")
+    node4 = Node("4")
+    node1 = node1.depends_on(node2)
+    node2 = node2.depends_on(node3)
+    node4 = node4.depends_on(node3)
+    node1 = node1.depends_on(node4)
+    node3 = node3.depends_on(node1)
+
+    to_remove = calculate_feedback_arc_set_on_graph([node1, node2, node3, node4])
+
+    assert set(to_remove) == {Edge(3, 1), Edge(4, 3)}
+
+
+def test_when_graph_is_acyclic_there_is_nothing_to_remove() -> None:
+    node1 = Node("1")
+    node2 = Node("2")
+    node3 = Node("3")
+    node4 = Node("4")
+    node1 = node1.depends_on(node2)
+    node2 = node2.depends_on(node3)
+    node1 = node1.depends_on(node4)
+
+    to_remove = calculate_feedback_arc_set_on_graph([node1, node2, node3, node4])
+
+    assert to_remove == []
diff --git a/tests/sorter/test_graph_topological_sort.py b/tests/sorter/test_graph_topological_sort.py
new file mode 100644
index 0000000..1af20ec
--- /dev/null
+++ b/tests/sorter/test_graph_topological_sort.py
@@ -0,0 +1,84 @@
+from smartschedule.sorter.graph_topological_sort import graph_topological_sort
+from smartschedule.sorter.node import Node
+from smartschedule.sorter.nodes import Nodes
+
+
+def test_topological_sort_with_simple_dependencies() -> None:
+    node_1 = Node("Node1")
+    node_2 = Node("Node2")
+    node_3 = Node("Node3")
+    node_4 = Node("Node4")
+
+    node_2 = node_2.depends_on(node_1)
+    node_3 = node_3.depends_on(node_1)
+    node_4 = node_4.depends_on(node_2)
+
+    nodes = Nodes([node_1, node_2, node_3, node_4])
+
+    sorted_nodes = graph_topological_sort(nodes)
+
+    assert len(sorted_nodes) == 3
+
+    assert node_1 in sorted_nodes[0]
+    assert node_2 in sorted_nodes[1]
+    assert node_3 in sorted_nodes[1]
+    assert node_4 in sorted_nodes[2]
+
+
+def test_topological_sort_with_linear_dependencies() -> None:
+    node_1 = Node("Node1")
+    node_2 = Node("Node2")
+    node_3 = Node("Node3")
+    node_4 = Node("Node4")
+    node_5 = Node("Node5")
+
+    node_2 = node_2.depends_on(node_1)
+    node_3 = node_3.depends_on(node_2)
+    node_4 = node_4.depends_on(node_3)
+    node_5 = node_5.depends_on(node_4)
+
+    nodes = Nodes([node_1, node_2, node_3, node_4, node_5])
+
+    sorted_nodes = graph_topological_sort(nodes)
+
+    assert len(sorted_nodes) == 5
+
+    assert node_5 in sorted_nodes[4]
+    assert len(sorted_nodes[4]) == 1
+
+    assert node_4 in sorted_nodes[3]
+    assert len(sorted_nodes[3]) == 1
+
+    assert node_3 in sorted_nodes[2]
+    assert len(sorted_nodes[2]) == 1
+
+    assert node_2 in sorted_nodes[1]
+    assert len(sorted_nodes[1]) == 1
+
+    assert node_1 in sorted_nodes[0]
+    assert len(sorted_nodes[0]) == 1
+
+
+def test_nodes_without_dependencies() -> None:
+    node_1 = Node("Node1")
+    node_2 = Node("Node2")
+
+    nodes = Nodes([node_1, node_2])
+
+    sorted_nodes = graph_topological_sort(nodes)
+
+    assert len(sorted_nodes) == 1
+
+
+def test_cyclic_dependency() -> None:
+    node_1 = Node("Node1")
+    node_2 = Node("Node2")
+
+    node_2 = node_2.depends_on(node_1)
+    node_1 = node_1.depends_on(node_2)
+
+    nodes = Nodes([node_1, node_2])
+
+    sorted_nodes = graph_topological_sort(nodes)
+
+    assert len(sorted_nodes) == 0
diff --git a/tests/test_parallelization.py b/tests/test_parallelization.py
deleted file mode 100644
index 401a100..0000000
--- a/tests/test_parallelization.py
+++ /dev/null
@@ -1,36 +0,0 @@
-from smartschedule.parallelization.stage import Stage
-from smartschedule.parallelization.stage_parallelization import StageParallelization
-
-
-def test_everything_can_be_done_in_parallel_when_there_are_no_deps() -> None:
-    stage_1 = Stage("stage_1")
-    stage_2 = Stage("stage_2")
-
-    sorted_stages = StageParallelization.of({stage_1, stage_2})
-
-    assert len(sorted_stages.all) == 1
-
-
-def test_simple_deps() -> None:
-    stage_1 = Stage("stage_1")
-    stage_2 = Stage("stage_2")
-    stage_3 = Stage("stage_3")
-    stage_4 = Stage("stage_4")
-    stage_2.depends_on(stage_1)
-    stage_3.depends_on(stage_1)
-    stage_4.depends_on(stage_2)
-
-    sorted_stages = StageParallelization.of({stage_1, stage_2, stage_3, stage_4})
-
-    assert str(sorted_stages) == "stage_1 | stage_2, stage_3 | stage_4"
-
-
-def test_cannot_be_done_when_there_is_a_cycle() -> None:
-    stage_1 = Stage("stage_1")
-    stage_2 = Stage("stage_2")
-    stage_2.depends_on(stage_1)
-    stage_1.depends_on(stage_2)  # making it cyclic
-
-    sorted_stages = StageParallelization.of({stage_1, stage_2})
-
-    assert len(sorted_stages.all) == 0