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_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_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_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_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)", ]