def test_repr(self):
arg = py.Symbol("a")
kwarg = py.Symbol("v")
assert (repr(py.FunctionCall(
"f", [arg],
{"k": kwarg
})) == f"FunctionCall('f', [{arg!r}], {{'k': {kwarg!r}}})")
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 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):
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_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 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 test_lines_with_comments(self, level: int):
x = py.Line.INDENT_UNIT
ob = py.Standalone(py.Symbol("a"), comments=["1", "2"])
assert [str(l) for l in ob.lines(level)] == [
x * level + "# 1",
x * level + "# 2",
x * level + "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_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_with_positional_and_kwargs(self):
assert (str(
py.FunctionCall(
"m.f",
[py.Literal(True),
py.FunctionCall("g", [py.Symbol("f")])],
{
"a": py.Literal(2),
"bc": py.Literal("x")
},
)) == "m.f(True, g(f), a=2, bc='x')")
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 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 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_repr(self):
assert repr(py.Symbol(" x'\" y ")) == "Symbol(' x\\'\" y ')"
def test_non_strings_raise_error(self):
with pytest.raises(TypeError):
assert str(py.Symbol(True))
def test_strings_appear_unchanged(self, s: str):
assert str(py.Symbol(s)) == s
def test_repr(self):
expr = py.Symbol("a")
assert (repr(py.Standalone(
expr,
comments=["hi"])) == f"Standalone({expr!r}, comments=['hi'])")
