Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 08:11:00 +00:00
Code Issues Projects Releases Wiki Activity

locust_task re-use locust_code.

Browse Source
This commit is contained in:
Zohar Lorberbaum 2016-03-27 21:42:52 -07:00
parent 9f77c80a32
commit ef3d24e8c8
2 changed files with 25 additions and 69 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

65
mitmproxy/flow_export.py
View File

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

29
test/mitmproxy/test_flow_export.py
View File

@ -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