def test_lines_with_nested_body(self, level: int):
x = py.Line.INDENT_UNIT
f = py.Function(
"func",
params=[],
statements=[
py.Assignment("a", py.Literal(2)),
py.IfElse(
[(py.Literal(True), [py.Assignment("b", py.Literal(3))])],
[
py.Assignment("b", py.Literal(4)),
py.Assignment("c", py.Literal(1)),
],
),
],
)
assert [str(l) for l in f.lines(level)] == [
x * level + "def func():",
x * (level + 1) + "a = 2",
x * (level + 1) + "if True:",
x * (level + 2) + "b = 3",
x * (level + 1) + "else:",
x * (level + 2) + "b = 4",
x * (level + 2) + "c = 1",
]
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_lines_with_hidden_comments(self, level: int):
x = py.Line.INDENT_UNIT
cond = py.Literal(True)
if_true = py.Assignment("x", py.Symbol("a"), comments=["tx", "ty"])
if_false = py.Assignment("x", py.Symbol("b"), comments=["fx", "fy"])
stmt = py.IfElse([(cond, [if_true])], [if_false], comments=["1", "2"])
assert [str(l) for l in stmt.lines(level, comments=False)] == [
x * level + "if True:",
*[str(l) for l in if_true.lines(level + 1, comments=False)],
x * level + "else:",
*[str(l) for l in if_false.lines(level + 1, comments=False)],
]
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",
]
def test_lines_for_if_else(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.Assignment("t", py.Literal(2))],
).lines(level)
] == [
x * level + "if t is None:",
x * (level + 1) + "t = 1",
x * level + "else:",
x * (level + 1) + "t = 2",
]
def test_init_with_no_condition_raises_error(self):
with pytest.raises(ValueError):
py.IfElse([])
with pytest.raises(ValueError):
py.IfElse([], else_block=[])
with pytest.raises(ValueError):
py.IfElse([], else_block=[py.Assignment("a", py.Literal(2))])
def test_repr(self):
cond = py.Literal(True)
if_true = [(cond, [py.Assignment("x", py.Symbol("a"))])]
stmt = py.IfElse(condition_blocks=if_true, comments=["hi"])
assert (
repr(stmt) ==
f"IfElse(condition_blocks={if_true!r}, else_block=None, comments=['hi'])"
)
def test_lines_with_comments(self, level: int):
x = py.Line.INDENT_UNIT
stmt = py.Assignment("x", py.Symbol("a"), comments=["1", "2"])
assert [str(l) for l in stmt.lines(level)] == [
x * level + "# 1",
x * level + "# 2",
x * level + "x = a",
]
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_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_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_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 from_task(cls, task: "Task") -> "Task2":
# TODO: Remove me as soon as the old Task is no longer used and Task2 is
# renamed to Task.
# See https://github.com/zalando-incubator/Transformer/issues/11.
t = cls(name=task.name, request=task.request)
if task.locust_request:
expr_view = py.ExpressionView(
name="this task's request field",
target=lambda: task.locust_request,
converter=lreq_to_expr,
)
else:
expr_view = py.ExpressionView(
name="this task's request field",
target=lambda: t.request,
converter=req_to_expr,
)
t.statements = [
*[py.OpaqueBlock(x) for x in task.locust_preprocessing],
py.Assignment("response", expr_view),
*[py.OpaqueBlock(x) for x in task.locust_postprocessing],
]
return t
def test_repr(self):
rhs = py.Symbol("a")
assert (repr(py.Assignment("x", rhs, comments=[
"hi"
])) == f"Assignment(lhs='x', rhs={rhs!r}, comments=['hi'])")
def test_simple(self, level: int):
x = py.Line.INDENT_UNIT
assert [
str(l) for l in py.Assignment("foo", py.Literal(3)).lines(level)
] == [x * level + "foo = 3"]
