forked from mozilla/amo-loadtest
-
Notifications
You must be signed in to change notification settings - Fork 0
/
locustfile.py
223 lines (193 loc) · 8.4 KB
/
locustfile.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
import logging
import os
import random
import re
import tempfile
import time
import uuid
from contextlib import contextmanager
from shutil import make_archive, rmtree
from zipfile import ZipFile
import lxml.html
from lxml.html import submit_form
from locust import HttpLocust, TaskSet, task
MAX_UPLOAD_POLL_ATTEMPTS = 200
ID_REGEX = re.compile('THIS_IS_THE_ID')
NAME_REGEX = re.compile('THIS_IS_THE_NAME')
root_path = os.path.dirname(__file__)
data_dir = os.path.join(root_path, 'data')
xpi_dir = os.path.join(root_path, 'add-ons')
xpis = [os.path.join(xpi_dir, xpi) for xpi in os.listdir(xpi_dir)]
log = logging.getLogger(__name__)
def get_random():
return str(uuid.uuid4())
def submit_url(step):
return '/en-US/developers/addon/submit/{step}'.format(step=step)
def get_xpi():
return uniqueify_xpi(random.choice(xpis))
@contextmanager
def uniqueify_xpi(path):
output_dir = tempfile.mkdtemp()
try:
xpi_dir = os.path.join(output_dir, 'xpi')
output_path = os.path.join(output_dir, 'addon')
xpi_name = os.path.basename(path)
xpi_path = os.path.join(output_dir, xpi_name)
with ZipFile(path) as original:
original.extractall(xpi_dir)
with open(os.path.join(xpi_dir, 'install.rdf')) as f:
install_rdf = f.read()
install_rdf = ID_REGEX.sub('{%s}' % get_random(), install_rdf)
install_rdf = NAME_REGEX.sub(get_random(), install_rdf)
with open(os.path.join(xpi_dir, 'install.rdf'), 'w') as f:
f.write(install_rdf)
archive_path = make_archive(output_path, 'zip', xpi_dir)
os.rename(archive_path, xpi_path)
with open(xpi_path) as f:
yield f
finally:
rmtree(output_dir)
class UserBehavior(TaskSet):
def on_start(self):
user_file = os.path.join(data_dir, 'loadtest-users.txt')
if not os.path.exists(user_file):
raise ValueError(
'User file does not exist: {}; did you generate one?'
.format(user_file))
users = []
with open(user_file, 'r') as f:
for line in f:
email, password = line.strip().split(':')
users.append({'email': email, 'password': password})
credentials = random.choice(users)
log.info('Running test with {}'.format(credentials))
self.login(credentials['email'], credentials['password'])
def submit_form(self, form=None, url=None, extra_values=None):
if form is None:
raise ValueError('form cannot be None; url={}'.format(url))
def submit(method, form_action_url, values):
values = dict(values)
if 'csrfmiddlewaretoken' not in values:
raise ValueError(
'Possibly the wrong form. Could not find '
'csrfmiddlewaretoken: {}'.format(repr(values)))
with self.client.post(
url or form_action_url, values,
allow_redirects=False, catch_response=True) as response:
if response.status_code not in (301, 302):
# This probably means the form failed and is displaying
# errors.
# TODO: scrape out the errors.
response.failure(
'Form submission did not redirect; status={}'
.format(response.status_code))
submit_form(form, open_http=submit, extra_values=extra_values)
def get_the_only_form_without_id(self, response_content):
"""
Gets the only form on the page that doesn't have an ID.
A lot of pages (login, registration) have a single form without an ID.
This is the one we want. The other forms on the page have IDs so we
can ignore them. I'm sure this will break one day.
"""
html = lxml.html.fromstring(response_content)
target_form = None
for form in html.forms:
if not form.attrib.get('id'):
target_form = form
if target_form is None:
raise ValueError(
'Could not find only one form without an ID; found: {}'
.format(html.forms))
return target_form
def login(self, email, password):
login_url = "/en-US/firefox/users/login"
resp = self.client.get(login_url)
login_form = self.get_the_only_form_without_id(resp.content)
self.submit_form(
form=login_form, url=login_url, extra_values={
"username": email,
"password": password})
def load_upload_form(self):
url = submit_url(2)
with self.client.get(
url, allow_redirects=False, catch_response=True) as response:
if response.status_code == 200:
html = lxml.html.fromstring(response.content)
return html.get_element_by_id('create-addon')
else:
more_info = ''
if response.status_code in (301, 302):
more_info = ('Location: {}'
.format(response.headers['Location']))
response.failure('Unexpected status: {}; {}'
.format(response.status_code, more_info))
def upload_addon(self, form):
url = submit_url(2)
csrfmiddlewaretoken = form.fields['csrfmiddlewaretoken']
with get_xpi() as addon_file:
with self.client.post(
'/en-US/developers/upload',
{'csrfmiddlewaretoken': csrfmiddlewaretoken},
files={'upload': addon_file},
name='/en-US/developers/upload ({})'.format(
os.path.basename(addon_file.name)),
allow_redirects=False,
catch_response=True) as response:
if response.status_code == 302:
poll_url = response.headers['location']
upload_uuid = self.poll_upload_until_ready(poll_url)
if upload_uuid:
form.fields['upload'] = upload_uuid
self.submit_form(form=form, url=url)
else:
response.failure('Unexpected status: {}'.format(
response.status_code))
@task(1)
def upload(self):
form = self.load_upload_form()
if form:
self.upload_addon(form)
@task(5)
def browse(self):
self.client.get('/en-US/firefox/')
self.client.get('/en-US/firefox/search/?q=pi&appver=45.0&platform=mac')
with self.client.get(
'/en-US/firefox/extensions/',
allow_redirects=False, catch_response=True) as response:
if response.status_code == 200:
html = lxml.html.fromstring(response.content)
addon_links = html.cssselect('.item.addon h3 a')
url = random.choice(addon_links).get('href')
self.client.get(url, name='/en-US/firefox/addon/:slug')
else:
response.failure('Unexpected status code {}'.format(
response.status_code))
def poll_upload_until_ready(self, url):
for i in xrange(MAX_UPLOAD_POLL_ATTEMPTS):
with self.client.get(url, allow_redirects=False,
name='/en-US/developers/upload/:uuid/json',
catch_response=True) as response:
try:
data = response.json()
except ValueError:
return response.failure(
'Failed to parse JSON when polling. '
'Status: {} content: {}'.format(
response.status_code, response.content))
if response.status_code == 200:
if data['error']:
return response.failure('Unexpected error: {}'.format(
data['error']))
elif data['validation']:
return data['upload']
else:
return response.failure('Unexpected status: {}'.format(
response.status_code))
time.sleep(1)
else:
response.failure('Upload did not complete in {} tries'.format(
MAX_UPLOAD_POLL_ATTEMPTS))
class WebsiteUser(HttpLocust):
task_set = UserBehavior
min_wait = 5000
max_wait = 9000