diff options
| -rw-r--r-- | mitmproxy/flow_export.py | 65 | ||||
| -rw-r--r-- | test/mitmproxy/test_flow_export.py | 29 |
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 |