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_empty_class(self, level: int):
c = py.Class("A", statements=[])
x = py.Line.INDENT_UNIT
assert [str(l) for l in c.lines(level)] == [
x * level + "class A:",
x * (level + 1) + "pass",
]
def locust_classes(scenarios: Sequence[Scenario]) -> List[py.Class]:
"""
Transforms scenarios into all Python classes needed by Locust (TaskSet and
Locust classes).
The only missing parts before a fully functional locustfile are:
- integrating all necessary set-up/tear-down statements:
- Python imports,
- apply global plugins,
- etc.
- serializing everything via transformer.python.
"""
classes = []
for scenario in scenarios:
taskset = locust_taskset(scenario)
locust_class = py.Class(
name=f"LocustFor{taskset.name}",
superclasses=["HttpLocust"],
statements=[
py.Assignment("task_set", py.Symbol(taskset.name)),
py.Assignment("weight", py.Literal(scenario.weight)),
py.Assignment("min_wait", py.Literal(LOCUST_MIN_WAIT_DELAY)),
py.Assignment("max_wait", py.Literal(LOCUST_MAX_WAIT_DELAY)),
],
)
classes.append(taskset)
classes.append(locust_class)
return classes
def test_repr(self):
stmts = [py.OpaqueBlock("raise")]
assert (
repr(
py.Class(name="C",
statements=stmts,
superclasses=["A"],
comments=["hi"])) ==
f"Class(name='C', statements={stmts!r}, superclasses=['A'], comments=['hi'])"
)
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_lines_with_hidden_comments(self, level: int):
stmt = py.OpaqueBlock("foo", comments=["x"])
c = py.Class("C",
statements=[stmt],
superclasses=[],
comments=["1", "2"])
x = py.Line.INDENT_UNIT
assert [str(l) for l in c.lines(level, comments=False)] == [
x * level + "class C:",
*[str(l) for l in stmt.lines(level + 1, comments=False)],
]
def test_class_with_superclasses(self, names: List[str]):
c = py.Class("A", statements=[], superclasses=names)
x = py.Line.INDENT_UNIT
if names:
expected = "(" + ", ".join(names) + ")"
else:
expected = ""
assert [str(l)
for l in c.lines()] == [f"class A{expected}:", x + "pass"]
def test_class_with_fields(self, level: int):
c = py.Class(
"A",
statements=[
py.Assignment("a", py.Literal(2)),
py.Assignment("b", py.Literal(3)),
],
)
x = py.Line.INDENT_UNIT
assert [str(l) for l in c.lines(level)] == [
x * level + "class A:",
x * (level + 1) + "a = 2",
x * (level + 1) + "b = 3",
]