1# Copyright 2016 Google Inc.
2#
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6import argparse
7import json
8import os
9import re
10import requests
11
12from os import listdir
13from os.path import isfile, join
14
15default_ops = [
16    "enable_gpu",
17    "post",
18    "info",
19    "cmd",
20    "img",
21    "batchList"
22]
23
24def Check(request):
25    assert(request.status_code == 200)
26    return request
27
28def WriteJson(request, path):
29    # Writes out pretty printed json
30    with open(path, 'wb+') as fd:
31        json.dump(request.json(), fd, sort_keys=True, indent=2,
32                  separators=(',', ': '))
33    return request
34
35def WritePng(request, path):
36    with open(path, 'wb+') as fd:
37        fd.write(request.content)
38
39
40# A simple class to drive testing
41class SkiaServeTester:
42    def __init__(self, url, output_dir):
43        self.url = url
44        self.output_dir = output_dir
45
46        # skp properties
47        self.skp = ''
48        self.skp_name = ''
49
50    def set_skp(self, skp_dir, skp_name):
51        self.skp = skp_dir + '/' + skp_name
52        self.skp_name = skp_name
53
54    def info(self):
55        return Check(requests.get(self.url + '/info'))
56
57    def post(self):
58        with open(self.skp, 'rb') as payload:
59            files = {'file': payload}
60
61            # upload skp
62            return Check(requests.post(self.url + '/new', files=files))
63
64    def cmd(self):
65        path = self.output_dir + '/' + self.skp_name + '.cmd.json'
66        return WriteJson(Check(requests.get(self.url + '/cmd')), path)
67
68    def img(self):
69        opcount = self.opcount()
70        url = self.url + '/img/' + str(opcount)
71        path = self.output_dir + '/' + self.skp_name + '.png'
72        return WritePng(Check(requests.get(url)), path)
73
74    def enable_gpu(self):
75        return Check(requests.post(self.url + '/enableGPU/1'))
76
77    def disable_gpu(self):
78        return Check(requests.post(self.url + '/enableGPU/0'))
79
80    def opcount(self):
81        r = self.cmd()
82        return len(r.json()['commands']) - 1 # why the minus 1 here?
83
84    def batchList(self):
85        path = self.output_dir + '/' + self.skp_name + '.batches.json'
86        return WriteJson(Check(requests.get(self.url + '/batches')), path)
87
88def main():
89    parser = argparse.ArgumentParser(description='Tester for SkiaServe')
90    parser.add_argument('--skp_dir', default='skps', type=str)
91    parser.add_argument('--url', default='http://localhost:8888', type=str)
92    parser.add_argument('--output_dir', default='results', type=str)
93    parser.add_argument('--match', default='.*', type=str)
94    parser.add_argument('--ops', nargs='+', default=default_ops)
95
96    args = parser.parse_args()
97    skp_dir = args.skp_dir
98    url = args.url
99    output_dir = args.output_dir
100    ops = args.ops
101
102    if not os.path.isdir(output_dir):
103        os.makedirs(output_dir)
104
105    skps = []
106    for skp in listdir(skp_dir):
107        if isfile(join(skp_dir, skp)) and re.match(args.match, skp):
108            skps.append(skp)
109
110    tester = SkiaServeTester(url, output_dir)
111
112    for skp_name in skps:
113        tester.set_skp(skp_dir, skp_name)
114        for op in ops:
115            getattr(tester, op)()
116
117if __name__ == "__main__":
118    main()
119