1810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt# Copyright 2016 Google Inc.
2810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt#
3810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt# Use of this source code is governed by a BSD-style license that can be
4810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt# found in the LICENSE file.
5810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
6810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualittimport argparse
7810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualittimport json
8810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualittimport os
9810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualittimport re
10810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualittimport requests
11810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
12810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualittfrom os import listdir
13810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualittfrom os.path import isfile, join
14810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
15beba90864ae4645b9f4f7a47065f567862ada12cjoshualittdefault_ops = [
16810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    "enable_gpu",
17810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    "post",
18810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    "info",
19810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    "cmd",
20beba90864ae4645b9f4f7a47065f567862ada12cjoshualitt    "img",
21beba90864ae4645b9f4f7a47065f567862ada12cjoshualitt    "batchList"
22810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt]
23810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
24810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualittdef Check(request):
25810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    assert(request.status_code == 200)
26810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    return request
27810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
28810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualittdef WriteJson(request, path):
29810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    # Writes out pretty printed json
30810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    with open(path, 'wb+') as fd:
31810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        json.dump(request.json(), fd, sort_keys=True, indent=2,
32810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt                  separators=(',', ': '))
33810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    return request
34810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
35810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualittdef WritePng(request, path):
36810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    with open(path, 'wb+') as fd:
37810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        fd.write(request.content)
38810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
39810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
40810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt# A simple class to drive testing
41810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualittclass SkiaServeTester:
42810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    def __init__(self, url, output_dir):
43810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        self.url = url
44810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        self.output_dir = output_dir
45810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
46810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        # skp properties
47810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        self.skp = ''
48810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        self.skp_name = ''
49810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
50810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    def set_skp(self, skp_dir, skp_name):
51810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        self.skp = skp_dir + '/' + skp_name
52810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        self.skp_name = skp_name
53810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
54810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    def info(self):
55810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        return Check(requests.get(self.url + '/info'))
56810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
57810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    def post(self):
58810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        with open(self.skp, 'rb') as payload:
59810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt            files = {'file': payload}
6063fd760a37905c45d26fc3d49cac261fad1b4808Ben Wagner
61810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt            # upload skp
62810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt            return Check(requests.post(self.url + '/new', files=files))
63810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
64810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    def cmd(self):
65810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        path = self.output_dir + '/' + self.skp_name + '.cmd.json'
66810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        return WriteJson(Check(requests.get(self.url + '/cmd')), path)
67810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
68810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    def img(self):
69810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        opcount = self.opcount()
70810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        url = self.url + '/img/' + str(opcount)
71810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        path = self.output_dir + '/' + self.skp_name + '.png'
72810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        return WritePng(Check(requests.get(url)), path)
73810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
74810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    def enable_gpu(self):
75810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        return Check(requests.post(self.url + '/enableGPU/1'))
7663fd760a37905c45d26fc3d49cac261fad1b4808Ben Wagner
77810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    def disable_gpu(self):
78810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        return Check(requests.post(self.url + '/enableGPU/0'))
79810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
80810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    def opcount(self):
81810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        r = self.cmd()
82810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        return len(r.json()['commands']) - 1 # why the minus 1 here?
83810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
84beba90864ae4645b9f4f7a47065f567862ada12cjoshualitt    def batchList(self):
85beba90864ae4645b9f4f7a47065f567862ada12cjoshualitt        path = self.output_dir + '/' + self.skp_name + '.batches.json'
86beba90864ae4645b9f4f7a47065f567862ada12cjoshualitt        return WriteJson(Check(requests.get(self.url + '/batches')), path)
87beba90864ae4645b9f4f7a47065f567862ada12cjoshualitt
88810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualittdef main():
89810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    parser = argparse.ArgumentParser(description='Tester for SkiaServe')
90810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    parser.add_argument('--skp_dir', default='skps', type=str)
91810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    parser.add_argument('--url', default='http://localhost:8888', type=str)
92810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    parser.add_argument('--output_dir', default='results', type=str)
93810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    parser.add_argument('--match', default='.*', type=str)
94beba90864ae4645b9f4f7a47065f567862ada12cjoshualitt    parser.add_argument('--ops', nargs='+', default=default_ops)
95810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
96810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    args = parser.parse_args()
97810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    skp_dir = args.skp_dir
98810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    url = args.url
99810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    output_dir = args.output_dir
100beba90864ae4645b9f4f7a47065f567862ada12cjoshualitt    ops = args.ops
101810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
102810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    if not os.path.isdir(output_dir):
103810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        os.makedirs(output_dir)
104810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
105810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    skps = []
106810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    for skp in listdir(skp_dir):
107810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        if isfile(join(skp_dir, skp)) and re.match(args.match, skp):
108810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt            skps.append(skp)
109810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
110810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    tester = SkiaServeTester(url, output_dir)
111810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
112810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    for skp_name in skps:
113810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        tester.set_skp(skp_dir, skp_name)
114810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt        for op in ops:
115810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt            getattr(tester, op)()
116810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt
117810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualittif __name__ == "__main__":
118810f77fb9c21a8f2e3373ff92ef302e22b585e4cjoshualitt    main()
119