locust_task re-use locust_code.

2 files changed, 25 insertions, 69 deletions

diff --git a/mitmproxy/flow_export.py b/mitmproxy/flow_export.py
index a9fd5cde..e2d46d93 100644
--- a/mitmproxy/flow_export.py
+++ b/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,10 +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)
+    file_name = "_".join(components)
+    name = re.sub('\W|^(?=\d)', '_', file_name)
     url = flow.request.scheme + "://" + flow.request.host + "/" + "/".join(components)
 
     args = ""
@@ -146,6 +148,7 @@ def locust_code(flow):
         args += "\n            data=data,"
 
     code = code.format(
+        name=name,
         url=url,
         headers=headers,
         params=params,
@@ -163,55 +166,9 @@ def locust_code(flow):
 
 
 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)
-    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,"
+    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]
 
-    data = ""
-    if flow.request.body:
-        data = "\n    data = '''%s'''\n" % flow.request.body
-        args += "\n        data=data,"
-
-    code = code.format(
-        name=name,
-        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) + '")
-
-    code = "\n".join("    " + i for i in code.splitlines())
-
-    return code
+    return task_code
diff --git a/test/mitmproxy/test_flow_export.py b/test/mitmproxy/test_flow_export.py
index 6c9a6756..d937135f 100644
--- a/test/mitmproxy/test_flow_export.py
+++ b/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 = {
@@ -312,13 +312,13 @@ class TestExportLocustTask():
             'header': 'qvalue',
             'content-length': '7',
         }
-    
+
         self.response = self.client.request(
             method='GET',
             url=url,
             headers=headers,
         )
-        """.strip()
+        """.strip() + '\n'
 
         assert flow_export.locust_task(flow) == result
 
@@ -330,14 +330,13 @@ class TestExportLocustTask():
         url = '' + self.locust.host + '/path'
         
         data = '''content'''
-    
+
         self.response = self.client.request(
             method='POST',
             url=url,
             data=data,
         )
-    
-        """.strip()
+        """.strip() + '\n'
 
         assert flow_export.locust_task(flow) == result
 
@@ -353,13 +352,13 @@ class TestExportLocustTask():
             'header': 'qvalue',
             'content-length': '7',
         }
-    
+
         params = {
             'query': 'param',
         }
-    
+
         data = '''content'''
-    
+
         self.response = self.client.request(
             method='PATCH',
             url=url,
@@ -367,6 +366,6 @@ class TestExportLocustTask():
             params=params,
             data=data,
         )
-        """.strip()
+        """.strip() + '\n'
 
         assert flow_export.locust_task(flow) == result