locust_task re-use locust_code.

mitmproxy/flow_export.py

@@ -103,10 +103,10 @@ def locust_code(flow):
         class UserBehavior(TaskSet):
             def on_start(self):
                 ''' on_start is called when a Locust start before any task is scheduled '''
-                self.flow()
+                self.{name}()
 
             @task()
-            def flow(self):
+            def {name}(self):
                 url = '{url}'
                 {headers}{params}{data}
                 self.response = self.client.request(
@@ -121,62 +121,12 @@ def locust_code(flow):
             task_set = UserBehavior
             min_wait = 1000
             max_wait = 3000
+""").strip()
 
-    """).strip()
-
-    components = map(lambda x: quote(x, safe=""), flow.request.path_components)
-    url = flow.request.scheme + "://" + flow.request.host + "/" + "/".join(components)
-
-    args = ""
-    headers = ""
-    if flow.request.headers:
-        lines = ["            '%s': '%s',\n" % (k, v) for k, v in flow.request.headers.fields if k.lower() not in ["host", "cookie"]]
-        headers += "\n        headers = {\n%s        }\n" % "".join(lines)
-        args += "\n            headers=headers,"
-
-    params = ""
-    if flow.request.query:
-        lines = ["            '%s': '%s',\n" % (k, v) for k, v in flow.request.query]
-        params = "\n        params = {\n%s        }\n" % "".join(lines)
-        args += "\n            params=params,"
-
-    data = ""
-    if flow.request.body:
-        data = "\n        data = '''%s'''\n" % flow.request.body
-        args += "\n            data=data,"
-
-    code = code.format(
-        url=url,
-        headers=headers,
-        params=params,
-        data=data,
-        method=flow.request.method,
-        args=args,
-    )
-
-    host = flow.request.scheme + "://" + flow.request.host
-    code = code.replace(host, "' + self.locust.host + '")
-    code = code.replace(quote_plus(host), "' + quote_plus(self.locust.host) + '")
-    code = code.replace(quote(host), "' + quote(self.locust.host) + '")
-
-    return code
-
-
-def locust_task(flow):
-    code = dedent("""
-    @task()
-    def {name}(self):
-        url = '{url}'
-        {headers}{params}{data}
-        self.response = self.client.request(
-            method='{method}',
-            url=url,{args}
-        )
-    """).strip()
 
     components = map(lambda x: quote(x, safe=""), flow.request.path_components)
     file_name = "_".join(components)
-    name = re.sub('\W|^(?=\d)','_', file_name)
+    name = re.sub('\W|^(?=\d)', '_', file_name)
     url = flow.request.scheme + "://" + flow.request.host + "/" + "/".join(components)
 
     args = ""
@@ -212,6 +162,13 @@ def locust_task(flow):
     code = code.replace(quote_plus(host), "' + quote_plus(self.locust.host) + '")
     code = code.replace(quote(host), "' + quote(self.locust.host) + '")
 
-    code = "\n".join("    " + i for i in code.splitlines())
-
     return code
+
+
+def locust_task(flow):
+    code = locust_code(flow)
+    start_task = len(code.split('@task')[0]) - 4
+    end_task = -19 - len(code.split('### Additional')[1])
+    task_code = code[start_task:end_task]
+
+    return task_code

test/mitmproxy/test_flow_export.py

@@ -188,10 +188,10 @@ from locust import HttpLocust, TaskSet, task
 class UserBehavior(TaskSet):
     def on_start(self):
         ''' on_start is called when a Locust start before any task is scheduled '''
-        self.flow()
+        self.path()
 
     @task()
-    def flow(self):
+    def path(self):
         url = '' + self.locust.host + '/path'
         
         headers = {
@@ -226,10 +226,10 @@ from locust import HttpLocust, TaskSet, task
 class UserBehavior(TaskSet):
     def on_start(self):
         ''' on_start is called when a Locust start before any task is scheduled '''
-        self.flow()
+        self.path()
 
     @task()
-    def flow(self):
+    def path(self):
         url = '' + self.locust.host + '/path'
         
         data = '''content'''
@@ -261,10 +261,10 @@ from locust import HttpLocust, TaskSet, task
 class UserBehavior(TaskSet):
     def on_start(self):
         ''' on_start is called when a Locust start before any task is scheduled '''
-        self.flow()
+        self.path()
 
     @task()
-    def flow(self):
+    def path(self):
         url = '' + self.locust.host + '/path'
         
         headers = {
@@ -318,7 +318,7 @@ class TestExportLocustTask():
             url=url,
             headers=headers,
         )
-        """.strip()
+        """.strip() + '\n'
 
         assert flow_export.locust_task(flow) == result
 
@@ -336,8 +336,7 @@ class TestExportLocustTask():
             url=url,
             data=data,
         )
-    
+        """.strip() + '\n'
-        """.strip()
 
         assert flow_export.locust_task(flow) == result
 
@@ -367,6 +366,6 @@ class TestExportLocustTask():
             params=params,
             data=data,
         )
-        """.strip()
+        """.strip() + '\n'
 
         assert flow_export.locust_task(flow) == result