yandex-tank/Tank/Plugins/Loadosophia.py
2014-01-16 15:09:38 +04:00

423 lines
17 KiB
Python

''' Module to have Loadosophia.org integration '''
from Tank.Plugins.Aggregator import AggregateResultListener, AggregatorPlugin
from Tank.Plugins.ApacheBenchmark import ApacheBenchmarkPlugin
from Tank.Plugins.Monitoring import MonitoringPlugin
from Tank.Plugins.Phantom import PhantomPlugin
from Tank.Plugins.WebOnline import WebOnlinePlugin
from tankcore import AbstractPlugin
from urllib2 import HTTPError
import StringIO
import cookielib
import gzip
import itertools
import json
import logging
import mimetools
import mimetypes
import os
import time
import urllib
import urllib2
class LoadosophiaPlugin(AbstractPlugin, AggregateResultListener):
''' Tank plugin with Loadosophia.org uploading '''
SECTION = 'loadosophia'
@staticmethod
def get_key():
return __file__
def __init__(self, core):
''' Constructor '''
AbstractPlugin.__init__(self, core)
self.loadosophia = LoadosophiaClient()
self.loadosophia.results_url = None
self.project_key = None
self.color = None
self.title = None
self.online_buffer = []
self.online_initiated = False
self.online_enabled = False
def get_available_options(self):
return ["token", "project", "test_title", "file_prefix", "color_flag", "online_enabled" ]
def configure(self):
self.loadosophia.address = self.get_option("address", "https://loadosophia.org/")
self.loadosophia.token = self.get_option("token", "")
self.loadosophia.file_prefix = self.get_option("file_prefix", "")
self.project_key = self.get_option("project", 'DEFAULT')
self.title = self.get_option("test_title", "")
self.color = self.get_option("color_flag", "")
if self.loadosophia.token:
self.online_enabled = int(self.get_option("online_enabled", "1"))
try:
aggregator = self.core.get_plugin_of_type(AggregatorPlugin)
aggregator.add_result_listener(self)
except KeyError:
self.log.debug("No aggregator for loadosophia")
def start_test(self):
if self.online_enabled:
try:
url=self.loadosophia.start_online(self.project_key, self.title)
self.log.info("Started active test: %s", url)
except Exception, exc:
self.log.warning("Problems starting online: %s", exc)
self.online_enabled = False
def aggregate_second(self, second_aggregate_data):
if self.online_enabled:
self.log.debug("Online buffer: %s", self.online_buffer)
self.online_buffer.append(second_aggregate_data)
if len(self.online_buffer) >= 5 or not self.online_initiated:
try:
self.loadosophia.send_online_data(self.online_buffer)
self.online_initiated = True
except Exception, exc:
self.log.warning("Problems sending online data: %s", exc)
self.online_buffer = []
def post_process(self, retcode):
if self.online_enabled:
if self.online_buffer:
try:
self.loadosophia.send_online_data(self.online_buffer)
except Exception, exc:
self.log.warning("Problems sending online data rests: %s", exc)
self.online_buffer = []
# mark test closed
try:
self.loadosophia.end_online()
except Exception, exc:
self.log.warning("Problems ending online: %s", exc)
main_file = None
# phantom
try:
phantom = self.core.get_plugin_of_type(PhantomPlugin)
if phantom.phantom:
main_file = phantom.phantom.phout_file
except KeyError:
self.log.debug("Phantom not found")
# ab
try:
apache_bench = self.core.get_plugin_of_type(ApacheBenchmarkPlugin)
main_file = apache_bench.out_file
except KeyError:
self.log.debug("AB not found")
# monitoring
mon_file = None
try:
mon = self.core.get_plugin_of_type(MonitoringPlugin)
mon_file = mon.data_file
except KeyError:
self.log.debug("Monitoring not found")
queue_id = self.loadosophia.send_results(self.project_key, main_file, [mon_file])
if self.title or self.color:
test_id = self.loadosophia.get_test_by_upload(queue_id)
if self.color:
self.loadosophia.set_color_flag(test_id, self.color)
if self.title:
self.loadosophia.set_test_title(test_id, self.title)
if queue_id:
self.log.info("Loadosophia.org upload succeeded, report link: %s", self.loadosophia.results_url)
try:
web = self.core.get_plugin_of_type(WebOnlinePlugin)
if not web.redirect:
web.redirect = self.loadosophia.results_url
time.sleep(1)
except KeyError:
self.log.debug("Web online not found")
return retcode
class LoadosophiaClient:
''' Loadosophia service client class '''
STATUS_DONE = 4
def __init__(self):
self.log = logging.getLogger(__name__)
self.token = None
self.address = None
self.file_prefix = ''
self.results_url = None
self.cookie_jar = cookielib.CookieJar()
def send_results(self, project, result_file, monitoring_files):
''' Send files to loadosophia '''
if not self.token:
self.log.warning("Loadosophia.org uploading disabled, please set loadosophia.token option to enable it, get token at https://loadosophia.org/service/upload/token/")
else:
if not self.address:
self.log.warning("Loadosophia.org uploading disabled, please set loadosophia.address option to enable it")
else:
self.log.info("Uploading to Loadosophia.org: %s %s %s", project, result_file, monitoring_files)
if not project:
self.log.info("Uploading to default project, please set loadosophia.project option to change this")
if not result_file or not os.path.exists(result_file) or not os.path.getsize(result_file):
self.log.warning("Empty results file, skip Loadosophia.org uploading: %s", result_file)
else:
return self.__send_checked_results(project, result_file, monitoring_files)
def __send_checked_results(self, project, result_file, monitoring_files):
''' internal wrapper to send request '''
# Create the form with simple fields
form = MultiPartForm()
form.add_field('projectKey', project)
form.add_field('token', self.token)
# Add main file
form.add_file_as_string('jtl_file', self.file_prefix + os.path.basename(result_file) + ".gz", self.__get_gzipped_file(result_file))
index = 0
for mon_file in monitoring_files:
if not mon_file or not os.path.exists(mon_file) or not os.path.getsize(mon_file):
self.log.warning("Skipped mon file: %s", mon_file)
continue
form.add_file_as_string('perfmon_' + str(index), self.file_prefix + os.path.basename(mon_file) + ".gz", self.__get_gzipped_file(mon_file))
index += 1
# Build the request
request = urllib2.Request(self.address + "api/file/upload/?format=json")
request.add_header('User-Agent', 'Yandex.Tank Loadosophia Uploader Module')
body = str(form)
request.add_header('Content-Type', form.get_content_type())
request.add_header('Content-Length', len(body))
request.add_data(body)
response = urllib2.urlopen(request)
if response.getcode() != 200:
self.log.debug("Full loadosophia.org response: %s", response.read())
raise RuntimeError("Loadosophia.org upload failed, response code %s instead of 200, see log for full response text" % response.getcode())
resp_str = response.read()
try:
res = json.loads(resp_str)
except Exception, exc:
self.log.debug("Failed to load json from str: %s", resp_str)
raise exc
self.results_url = self.address + 'api/file/status/' + res[1][0]['QueueID'] + '/?redirect=true'
return res[1][0]['QueueID']
def __get_gzipped_file(self, result_file):
''' gzip file '''
out = StringIO.StringIO()
fhandle = gzip.GzipFile(fileobj=out, mode='w')
fhandle.write(open(result_file, 'r').read())
fhandle.close()
return out.getvalue()
def get_test_by_upload(self, queue_id):
self.log.info("Waiting for Loadosophia.org to process file...")
while True:
time.sleep(1)
status = self.get_upload_status(queue_id)
if status['UserError']:
raise HTTPError("Loadosophia processing error: " + status['UserError'])
if int(status['status']) == self.STATUS_DONE:
self.results_url = self.address + 'gui/' + status['TestID'] + '/'
return status['TestID']
def get_upload_status(self, queue_id):
self.log.debug("Requesting file status: %s", queue_id)
form = MultiPartForm()
form.add_field('token', self.token)
request = urllib2.Request(self.address + "api/file/status/" + queue_id + "/?format=json")
request.add_header('User-Agent', 'Yandex.Tank Loadosophia Uploader Module')
body = str(form)
request.add_header('Content-Type', form.get_content_type())
request.add_header('Content-Length', len(body))
request.add_data(body)
response = urllib2.urlopen(request)
if response.getcode() != 200:
self.log.debug("Full loadosophia.org response: %s", response.read())
raise RuntimeError("Loadosophia.org request failed, response code %s instead of 200, see log for full response text" % response.getcode())
res = json.loads(response.read())
self.log.debug("Status info: %s", res)
return res[1][0]
def set_color_flag(self, test_id, color):
form = MultiPartForm()
form.add_field('token', self.token)
request = urllib2.Request(self.address + "api/test/edit/color/" + test_id + "/?format=json&color=" + color)
request.add_header('User-Agent', 'Yandex.Tank Loadosophia Uploader Module')
body = str(form)
request.add_header('Content-Type', form.get_content_type())
request.add_header('Content-Length', len(body))
request.add_data(body)
response = urllib2.urlopen(request)
if response.getcode() != 204:
self.log.debug("Full loadosophia.org response: %s", response.read())
raise RuntimeError("Loadosophia.org request failed, response code %s instead of 204, see log for full response text" % response.getcode())
def set_test_title(self, test_id, title):
self.log.debug("Set test title: %s", title)
form = MultiPartForm()
form.add_field('token', self.token)
request = urllib2.Request(self.address + "api/test/edit/title/" + test_id + "/?format=json&" + urllib.urlencode({"title": title}))
request.add_header('User-Agent', 'Yandex.Tank Loadosophia Uploader Module')
body = str(form)
request.add_header('Content-Type', form.get_content_type())
request.add_header('Content-Length', len(body))
request.add_data(body)
response = urllib2.urlopen(request)
if response.getcode() != 204:
self.log.debug("Full loadosophia.org response: %s", response.read())
raise RuntimeError("Loadosophia.org request failed, response code %s instead of 204, see log for full response text" % response.getcode())
def start_online(self, project, title):
self.log.info("Initiating Loadosophia.org active test...")
data = urllib.urlencode({'projectKey': project, 'token':self.token, 'title': title})
opener = urllib2.build_opener(urllib2.HTTPCookieProcessor(self.cookie_jar))
url = self.address + "api/active/receiver/start/"
response = opener.open(url, data)
if response.getcode() != 201:
self.log.warn("Failed to start active test: %s", response.getcode())
self.log.debug("Failed to start active test: %s", response.read())
self.cookie_jar.clear_session_cookies()
online_id=json.loads(response.read())
return self.address + "gui/active/"+online_id['OnlineID']+'/'
def end_online(self):
self.log.debug("Ending Loadosophia online test")
opener = urllib2.build_opener(urllib2.HTTPCookieProcessor(self.cookie_jar))
url = self.address + "api/active/receiver/stop/"
response = opener.open(url)
if response.getcode() != 205:
self.log.warn("Failed to end active test: %s", response.getcode())
self.log.debug("Failed to end active test: %s", response.read())
self.cookie_jar.clear_session_cookies()
def send_online_data(self, data_buffer):
data = []
for sec in data_buffer:
item = sec.overall
json_item = {
"ts": str(sec.time),
"threads": item.active_threads,
"rps": item.RPS,
"planned_rps": item.planned_requests,
"avg_rt": item.avg_response_time,
"quantiles": item.quantiles,
"rc": item.http_codes,
"net": item.net_codes
}
data.append(json_item)
self.log.debug("Sending online data: %s", json.dumps(data))
data_str = urllib.urlencode({'data': json.dumps(data)})
opener = urllib2.build_opener(urllib2.HTTPCookieProcessor(self.cookie_jar))
url = self.address + "api/active/receiver/data/"
response = opener.open(url, data_str)
if response.getcode() != 202:
self.log.warn("Failed to push data: %s", response.getcode())
# =================================================================
class MultiPartForm(object):
"""Accumulate the data to be used when posting a form.
http://blog.doughellmann.com/2009/07/pymotw-urllib2-library-for-opening-urls.html """
def __init__(self):
self.form_fields = []
self.files = []
self.boundary = mimetools.choose_boundary()
return
def get_content_type(self):
''' returns content type '''
return 'multipart/form-data; boundary=%s' % self.boundary
def add_field(self, name, value):
"""Add a simple field to the form data."""
self.form_fields.append((name, value))
return
def add_file_as_string(self, fieldname, filename, body, mimetype=None):
''' add raw string file '''
if mimetype is None:
mimetype = mimetypes.guess_type(filename)[0] or 'application/octet-stream'
self.files.append((fieldname, filename, mimetype, body))
return
def add_file(self, fieldname, filename, file_handle, mimetype=None):
"""Add a file to be uploaded."""
body = file_handle.read()
self.add_file_as_string(fieldname, filename, body, mimetype)
return
def __str__(self):
"""Return a string representing the form data, including attached files."""
# Build a list of lists, each containing "lines" of the
# request. Each part is separated by a boundary string.
# Once the list is built, return a string where each
# line is separated by '\r\n'.
parts = []
part_boundary = '--' + self.boundary
# Add the form fields
parts.extend(
[ part_boundary,
'Content-Disposition: form-data; name="%s"' % name,
'',
value,
]
for name, value in self.form_fields
)
# Add the files to upload
parts.extend(
[ part_boundary,
'Content-Disposition: file; name="%s"; filename="%s"' % \
(field_name, filename),
'Content-Type: %s' % content_type,
'',
body,
]
for field_name, filename, content_type, body in self.files
)
# Flatten the list and add closing boundary marker,
# then return CR+LF separated data
flattened = list(itertools.chain(*parts))
flattened.append('--' + self.boundary + '--')
flattened.append('')
return '\r\n'.join(flattened)