From 76b351684ea625ae72a75bbc7e35b66e138e9101 Mon Sep 17 00:00:00 2001
From: Lars Holmberg <lars.holmberg@svenskaspel.se>
Date: Sun, 17 Dec 2023 21:17:42 +0100
Subject: [PATCH] update plugin example

---
 har2locust/extra_plugins/plugin_example.py | 33 ++++++----
 tests/outputs/login_plugin.py              | 70 +++++++++++-----------
 2 files changed, 55 insertions(+), 48 deletions(-)

diff --git a/har2locust/extra_plugins/plugin_example.py b/har2locust/extra_plugins/plugin_example.py
index 99447c1..459632c 100644
--- a/har2locust/extra_plugins/plugin_example.py
+++ b/har2locust/extra_plugins/plugin_example.py
@@ -35,7 +35,15 @@ def visit_FunctionDef(self, node: FunctionDef) -> FunctionDef:
             for i in range(len(node.body)):
                 try:
                     if node.body[i].items[0].context_expr.args[1].value == "/player/1/authenticate/testlogin":  # type: ignore
-                        node.body[i] = parse("self.auth()").body[0]
+                        block = parse(
+                            """
+self.customer = next(self.customer_iterator)
+self.auth()
+"""
+                        )
+                        node.body[i] = block.body[0]
+                        # yea, the next line modifies what we're iterating over so we'll miss the last element, which is ugly
+                        node.body.insert(i + 1, block.body[1])
                     # if node.body[i].items[0].context_expr.args[1].value == "/wager/9/wagers":  # type: ignore
                     #     json_param = [
                     #         kw_arg.value
@@ -51,17 +59,18 @@ def visit_FunctionDef(self, node: FunctionDef) -> FunctionDef:
                 except:
                     pass
 
-            # wrap the entire task function body in a with-block.
-            if node.name == "t":
-                with_block = parse(
-                    f"""
-with self.reader.user() as self.customer:
-    pass
-                    """
-                ).body[0]
-                assert isinstance(with_block, With), with_block
-                with_block.body = node.body
-                node.body = [with_block]
+            # Old school
+            #             # wrap the entire task function body in a with-block.
+            #             if node.name == "t":
+            #                 with_block = parse(
+            #                     f"""
+            # with self.reader.user() as self.customer:
+            #     pass
+            #                     """
+            #                 ).body[0]
+            #                 assert isinstance(with_block, With), with_block
+            #                 with_block.body = node.body
+            #                 node.body = [with_block]
             self.generic_visit(node)
             return node
 
diff --git a/tests/outputs/login_plugin.py b/tests/outputs/login_plugin.py
index 9a333d6..3a96c33 100644
--- a/tests/outputs/login_plugin.py
+++ b/tests/outputs/login_plugin.py
@@ -13,42 +13,40 @@ class login(RestUser):
 
     @task
     def t(self):
-        with self.reader.user() as self.customer:
-            self.auth()
-            with self.rest_(
-                "GET",
-                "/player/1/customizedsettings",
-                headers={"accept": "application/json, text/javascript, */*; q=0.01"},
-            ) as resp:
-                pass
-            with self.client.request(
-                "GET",
-                "https://spela.test3.svenskaspel.se/",
-                headers={
-                    "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9"
-                },
-                catch_response=True,
-            ) as resp:
-                pass
-            with self.client.request(
-                "GET",
-                "https://spela.test3.svenskaspel.se/logga-in/uppdaterade-villkor?returnUrl=%2F",
-                headers={
-                    "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9"
-                },
-                catch_response=True,
-            ) as resp:
-                pass
-            with self.rest(
-                "POST", "/player/1/terms", headers={"accept": "application/json, text/javascript, */*; q=0.01"}, json={}
-            ) as resp:
-                pass
-            with self.rest_(
-                "GET",
-                "/player/1/info?include=accountBalance",
-                headers={"accept": "application/json, text/javascript, */*; q=0.01"},
-            ) as resp:
-                pass
+        self.customer = next(self.customer_iterator)
+        self.auth()
+        with self.rest_(
+            "GET", "/player/1/customizedsettings", headers={"accept": "application/json, text/javascript, */*; q=0.01"}
+        ) as resp:
+            pass
+        with self.client.request(
+            "GET",
+            "https://spela.test3.svenskaspel.se/",
+            headers={
+                "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9"
+            },
+            catch_response=True,
+        ) as resp:
+            pass
+        with self.client.request(
+            "GET",
+            "https://spela.test3.svenskaspel.se/logga-in/uppdaterade-villkor?returnUrl=%2F",
+            headers={
+                "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9"
+            },
+            catch_response=True,
+        ) as resp:
+            pass
+        with self.rest(
+            "POST", "/player/1/terms", headers={"accept": "application/json, text/javascript, */*; q=0.01"}, json={}
+        ) as resp:
+            pass
+        with self.rest_(
+            "GET",
+            "/player/1/info?include=accountBalance",
+            headers={"accept": "application/json, text/javascript, */*; q=0.01"},
+        ) as resp:
+            pass
 
 
 if __name__ == "__main__":