def test_nested(self):
assert (str(
py.BinaryOp(
py.Literal(2),
"+",
py.BinaryOp(
py.BinaryOp(py.Literal(3), "-", py.Literal(4)),
"*",
py.Literal(5),
),
)) == "2 + ((3 - 4) * 5)")
def test_lines_for_single_if(self, level: int):
x = py.Line.INDENT_UNIT
assert [
str(l) for l in py.IfElse([(
py.BinaryOp(py.Symbol("t"), "is", py.Literal(None)),
[py.Assignment("t", py.Literal(1))],
)]).lines(level)
] == [x * level + "if t is None:", x * (level + 1) + "t = 1"]
def locust_imports() -> py.Program:
is_post_1 = py.BinaryOp(py.Symbol("LOCUST_MAJOR_VERSION"), ">=", py.Literal(1),)
imports_pre_1 = [
py.Import(
["HttpLocust", "TaskSequence", "TaskSet", "seq_task", "task"],
source="locust",
)
]
imports_post_1 = [
py.Import(
["HttpUser", "SequentialTaskSet", "TaskSet", "task"], source="locust",
),
py.Assignment("HttpLocust", py.Symbol("HttpUser")),
py.Assignment("TaskSequence", py.Symbol("SequentialTaskSet")),
py.Function("seq_task", ["_"], [py.Return(py.Symbol("task"))]),
]
return [
py.IfElse([(is_post_1, imports_post_1)], imports_pre_1),
]
def test_lines_for_if_elif_else_with_no_statements(self, level: int):
x = py.Line.INDENT_UNIT
assert [
str(l) for l in py.IfElse(
[
(py.BinaryOp(py.Symbol("t"), "is", py.Literal(None)), []),
(py.Literal(False), []),
(py.Literal(True), []),
],
[],
).lines(level)
] == [
x * level + "if t is None:",
x * (level + 1) + "pass",
x * level + "elif False:",
x * (level + 1) + "pass",
x * level + "elif True:",
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)
is_post_1 = py.BinaryOp(py.Symbol("LOCUST_MAJOR_VERSION"), ">=", py.Literal(1))
tasks = py.IfElse(
[
(
is_post_1,
[py.Assignment("tasks", py.Literal([py.Symbol(taskset.name)]))],
)
],
[py.Assignment("task_set", py.Symbol(taskset.name))],
)
locust_class = py.Class(
name=f"LocustFor{taskset.name}",
superclasses=["HttpLocust"],
statements=[
tasks,
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_lines_for_if_elif(self, level: int):
x = py.Line.INDENT_UNIT
assert [
str(l) for l in py.IfElse([
(
py.BinaryOp(py.Symbol("t"), "is", py.Literal(None)),
[py.Assignment("t", py.Literal(1))],
),
(
py.Literal(False),
[
py.Assignment("t", py.Literal(2)),
py.Standalone(py.FunctionCall("f", [py.Symbol("t")])),
],
),
]).lines(level)
] == [
x * level + "if t is None:",
x * (level + 1) + "t = 1",
x * level + "elif False:",
x * (level + 1) + "t = 2",
x * (level + 1) + "f(t)",
]
def test_simple(self):
assert str(py.BinaryOp(py.Literal(2), "**",
py.Literal(10))) == "2 ** 10"
