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 test_lines_with_no_params_and_no_body(self, level: int):
f = py.Function("f", params=[], statements=[])
x = py.Line.INDENT_UNIT
assert [str(l) for l in f.lines(level)] == [
x * level + "def f():",
x * (level + 1) + "pass",
]
def test_lines_with_hidden_comments(self, level: int):
stmt = py.OpaqueBlock("foo", comments=["x"])
f = py.Function("f", params=[], statements=[stmt], comments=["1", "2"])
x = py.Line.INDENT_UNIT
assert [str(l) for l in f.lines(level, comments=False)] == [
x * level + "def f():",
*[str(l) for l in stmt.lines(level + 1, comments=False)],
]
def test_lines_with_complex_params(self):
f = py.Function("f",
params=["x: int", "abc: bool = True", "*z: str"],
statements=[])
assert [str(l) for l in f.lines()] == [
"def f(x: int, abc: bool = True, *z: str):",
py.Line.INDENT_UNIT + "pass",
]
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_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):
stmts = [py.OpaqueBlock("raise")]
assert (
repr(
py.Function(name="f",
params=["a"],
statements=stmts,
comments=["hi"])) ==
f"Function(name='f', params=['a'], statements={stmts!r}, comments=['hi'])"
)
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 _locust_task(task: Union[Task, Task2]) -> py.Function:
"""
Transforms a Task into the Python code expected by Locust.
This function is private because it does not return a complete Locust task
(the @task decorator is missing) and should therefore not be used for that
purpose by unsuspecting users.
"""
if isinstance(task, Task):
# TODO: remove when Task2 has replaced Task.
# See https://github.com/zalando-incubator/Transformer/issues/11.
task = Task2.from_task(task)
return py.Function(name=task.name, params=["self"], statements=task.statements)
def test_lines_with_simple_body(self, level: int):
f = py.Function(
"f",
params=[],
statements=[
py.OpaqueBlock("print('Hello!')"),
py.OpaqueBlock("return")
],
)
x = py.Line.INDENT_UNIT
assert [str(l) for l in f.lines(level)] == [
x * level + "def f():",
x * (level + 1) + "print('Hello!')",
x * (level + 1) + "return",
]
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_repr(self):
f = py.Function("f", params=[], statements=[])
assert (repr(py.Decoration(
"task", f,
comments=["hi"])) == f"Decoration('task', {f!r}, comments=['hi'])")
def test_lines_with_simple_params(self):
f = py.Function("f", params=["x", "y"], statements=[])
assert [str(l) for l in f.lines()] == [
"def f(x, y):",
py.Line.INDENT_UNIT + "pass",
]
