def locust_taskset(scenario: Scenario) -> py.Class:
"""
Transforms a scenario (potentially containing other scenarios) into a Locust
TaskSet definition.
"""
if any(isinstance(child, (Task, Task2)) for child in scenario.children):
ts_type = TaskSetType.Sequence
else:
ts_type = TaskSetType.Set
fields: List[py.Statement] = []
for i, child in enumerate(scenario.children, start=1):
seq_decorator = f"seq_task({i})"
if isinstance(child, (Task2, Task)):
fields.append(py.Decoration(seq_decorator, _locust_task(child)))
elif isinstance(child, Scenario):
field = py.Decoration(f"task({child.weight})", locust_taskset(child))
if ts_type is TaskSetType.Sequence:
field = py.Decoration(seq_decorator, field)
fields.append(field)
else:
wrong_type = child.__class__.__qualname__
scenario_type = scenario.__class__.__qualname__
raise TypeError(
f"unexpected type {wrong_type} in {scenario_type}.children: {child!r}"
)
return py.Class(scenario.name, superclasses=[str(ts_type.value)], statements=fields)
def test_nested_decorators(self, level: int):
f = py.Function("f",
params=[],
statements=[py.Assignment("a", py.Symbol("f"))])
first = py.Decoration("task(2)", f)
second = py.Decoration("task_seq(1)", first)
x = py.Line.INDENT_UNIT
assert [str(l) for l in second.lines(level)] == [
x * level + "@task_seq(1)",
x * level + "@task(2)",
*[str(l) for l in f.lines(level)],
]
def test_lines_with_hidden_comments(self, level: int):
f = py.Function("f", params=[], statements=[], comments=["1", "2"])
d = py.Decoration("task", f, comments=["x", "y"])
x = py.Line.INDENT_UNIT
assert [str(l) for l in d.lines(level, comments=False)] == [
x * level + "@task",
x * level + "def f():",
x * (level + 1) + "pass",
]
def test_with_a_class(self, level: int):
c = py.Class("C",
superclasses=[],
statements=[py.Assignment("a: int", py.Literal(1))])
d = py.Decoration("task", c)
x = py.Line.INDENT_UNIT
assert [str(l) for l in d.lines(level)] == [
x * level + "@task",
*[str(l) for l in c.lines(level)],
]
def test_with_a_function(self, level: int):
f = py.Function("f",
params=[],
statements=[py.Assignment("a", py.Symbol("f"))])
d = py.Decoration("task(2)", f)
x = py.Line.INDENT_UNIT
assert [str(l) for l in d.lines(level)] == [
x * level + "@task(2)",
*[str(l) for l in f.lines(level)],
]
def test_repr(self):
f = py.Function("f", params=[], statements=[])
assert (repr(py.Decoration(
"task", f,
comments=["hi"])) == f"Decoration('task', {f!r}, comments=['hi'])")