File size: 7,412 Bytes
a164e13 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 |
from typing import Any, Dict, Optional
import networkx
import pytest
from mergekit.common import ImmutableMap
from mergekit.graph import Executor, Task
EXECUTION_COUNTS: Dict[Task, int] = {}
class DummyTask(Task):
result: Any
dependencies: ImmutableMap[str, Task]
name: str = "DummyTask"
grouplabel: Optional[str] = None
execution_count: int = 0
def arguments(self):
return self.dependencies
def group_label(self) -> Optional[str]:
return self.grouplabel
def execute(self, **kwargs):
EXECUTION_COUNTS[self] = EXECUTION_COUNTS.get(self, 0) + 1
return self.result
def create_mock_task(name, result=None, dependencies=None, group_label=None):
if dependencies is None:
dependencies = {}
return DummyTask(
result=result,
dependencies=ImmutableMap(data=dependencies),
name=name,
grouplabel=group_label,
)
# Test cases for the Task implementation
class TestTaskClass:
def test_task_execute(self):
# Testing the execute method
task = create_mock_task("task1", result=42)
assert task.execute() == 42, "Task execution did not return expected result"
def test_task_priority(self):
task = create_mock_task("task1")
assert task.priority() == 0, "Default priority should be 0"
def test_task_group_label(self):
task = create_mock_task("task1")
assert task.group_label() is None, "Default group label should be None"
# Test cases for the Executor implementation
class TestExecutorClass:
def test_executor_initialization(self):
# Testing initialization with single task
task = create_mock_task("task1")
executor = Executor([task])
assert executor.targets == [
task
], "Executor did not initialize with correct targets"
def test_executor_empty_list(self):
list(Executor([]).run())
def test_executor_scheduling(self):
# Testing scheduling with dependencies
task1 = create_mock_task("task1", result=1)
task2 = create_mock_task("task2", result=2, dependencies={"task1": task1})
executor = Executor([task2])
assert (
len(executor._make_schedule([task2])) == 2
), "Schedule should include two tasks"
def test_executor_dependency_building(self):
# Testing dependency building
task1 = create_mock_task("task1")
task2 = create_mock_task("task2", dependencies={"task1": task1})
executor = Executor([task2])
dependencies = executor._build_dependencies([task2])
assert task1 in dependencies[task2], "Task1 should be a dependency of Task2"
def test_executor_run(self):
# Testing execution through the run method
task1 = create_mock_task("task1", result=10)
task2 = create_mock_task("task2", result=20, dependencies={"task1": task1})
executor = Executor([task2])
results = list(executor.run())
assert (
len(results) == 1 and results[0][1] == 20
), "Executor run did not yield correct results"
def test_executor_execute(self):
# Testing execute method for side effects
task1 = create_mock_task("task1", result=10)
executor = Executor([task1])
# No assert needed; we're ensuring no exceptions are raised and method completes
executor.execute()
def test_dependency_ordering(self):
# Testing the order of task execution respects dependencies
task1 = create_mock_task("task1", result=1)
task2 = create_mock_task("task2", result=2, dependencies={"task1": task1})
task3 = create_mock_task("task3", result=3, dependencies={"task2": task2})
executor = Executor([task3])
schedule = executor._make_schedule([task3])
assert schedule.index(task1) < schedule.index(
task2
), "Task1 should be scheduled before Task2"
assert schedule.index(task2) < schedule.index(
task3
), "Task2 should be scheduled before Task3"
class TestExecutorGroupLabel:
def test_group_label_scheduling(self):
# Create tasks with group labels and dependencies
task1 = create_mock_task("task1", group_label="group1")
task2 = create_mock_task(
"task2", dependencies={"task1": task1}, group_label="group1"
)
task3 = create_mock_task("task3", group_label="group2")
task4 = create_mock_task(
"task4", dependencies={"task2": task2, "task3": task3}, group_label="group1"
)
# Initialize Executor with the tasks
executor = Executor([task4])
# Get the scheduled tasks
schedule = executor._make_schedule([task4])
# Check if tasks with the same group label are scheduled consecutively when possible
group_labels_in_order = [
task.group_label() for task in schedule if task.group_label()
]
assert group_labels_in_order == [
"group1",
"group1",
"group2",
"group1",
], "Tasks with same group label are not scheduled consecutively"
def test_group_label_with_dependencies(self):
# Creating tasks with dependencies and group labels
task1 = create_mock_task("task1", result=1, group_label="group1")
task2 = create_mock_task(
"task2", result=2, dependencies={"task1": task1}, group_label="group2"
)
task3 = create_mock_task(
"task3", result=3, dependencies={"task2": task2}, group_label="group1"
)
executor = Executor([task3])
schedule = executor._make_schedule([task3])
scheduled_labels = [
task.group_label() for task in schedule if task.group_label()
]
# Check if task3 is scheduled after task1 and task2 due to dependency, even though it has the same group label as task1
group1_indices = [
i for i, label in enumerate(scheduled_labels) if label == "group1"
]
group2_index = scheduled_labels.index("group2")
assert (
group1_indices[-1] > group2_index
), "Task with the same group label but later dependency was not scheduled after different group label"
class TestExecutorSingleExecution:
def test_single_execution_per_task(self):
EXECUTION_COUNTS.clear()
shared_task = create_mock_task("shared_task", result=100)
task1 = create_mock_task("task1", dependencies={"shared": shared_task})
task2 = create_mock_task("task2", dependencies={"shared": shared_task})
task3 = create_mock_task("task3", dependencies={"task1": task1, "task2": task2})
Executor([task3]).execute()
assert shared_task in EXECUTION_COUNTS, "Dependency not executed"
assert (
EXECUTION_COUNTS[shared_task] == 1
), "Shared dependency should be executed exactly once"
class CircularTask(Task):
def arguments(self) -> Dict[str, Task]:
return {"its_a_me": self}
def execute(self, **_kwargs) -> Any:
assert False, "Task with circular dependency executed"
class TestExecutorCircularDependency:
def test_circular_dependency(self):
with pytest.raises(networkx.NetworkXUnfeasible):
Executor([CircularTask()]).execute()
if __name__ == "__main__":
pytest.main()
|